diff --git a/guided_upsampling.py b/guided_upsampling.py new file mode 100644 index 0000000..5147067 --- /dev/null +++ b/guided_upsampling.py @@ -0,0 +1,152 @@ +# Copyright (c) 2021 Lightricks. All rights reserved. + +import numpy as np +import cv2 + + +def bias(x, b=0.8): + denom = ((1 / b) - 2) * (1 - x) + 1 + return x / denom + + +def probability_to_confidence(probabilty, low_thresh=0.3, high_thresh=0.5): + eps = 0.01 + + low = probabilty < low_thresh + high = probabilty > high_thresh + confidence_low = bias((low_thresh - probabilty[low]) / low_thresh) + confidence_high = bias((probabilty[high] - high_thresh) / (1 - high_thresh)) + confidence = np.zeros_like(probabilty) + confidence[low] = confidence_low + confidence[high] = confidence_high + confidence = np.maximum(eps, confidence) + return confidence + + +def downsample2_antialiased(X): + kernel = np.array([1, 3, 3, 1]) / 8 + dst = cv2.sepFilter2D(X, -1, kernel, kernel, anchor=(1, 1), borderType=cv2.BORDER_REPLICATE) + return dst[::2, ::2] + + +def resize_helper(X, shape): + X = X.squeeze() + while np.all(np.array(X.shape[:2]) >= np.array(shape) * 2): + X = downsample2_antialiased(X) + return cv2.resize(X, dsize=tuple(shape[1::-1]), interpolation=cv2.INTER_LINEAR) + + +def resize(X, shape): + if X.ndim == 2 or X.shape[2] <= 4: + return resize_helper(X, shape) + # opencv doesn't work on more than 4 channels + X1 = resize_helper(X[..., :3], shape) + X2 = resize_helper(X[..., 3:], shape) + return np.concatenate([X1, X2], axis=2) + + +def outer_product_images(X, Y): + assert X.shape[-1] == 3 and Y.shape[-1] == 3 + X_flat = X[..., :, np.newaxis] + Y_flat = Y[..., np.newaxis, :] + + outer = np.matmul(X_flat, Y_flat) + ind = np.triu_indices(3) + outer = outer[..., ind[0], ind[1]] + return outer.reshape(X.shape[:-1] + (6,)) + + +def smooth_upsample(X, size, num_steps=None): + if num_steps is None: + log4ratio = np.max(0.5 * np.log2(np.array(size) / X.shape[:2])) + num_steps = np.maximum(1, log4ratio.round().astype(np.int)) + ratio = np.array(size) / X.shape[:2] + ratio_per_step = np.array(X.shape[:2]) * ratio / num_steps + for step in np.arange(1, num_steps + 1): + target_shape_for_step = np.round(step * ratio_per_step).astype(np.int) + X = resize(X, target_shape_for_step) + return X + + +def solve_image_ldl3(A, b): + A11, A12, A13, A22, A23, A33 = np.split(A, A.shape[-1], axis=-1) + b1, b2, b3 = np.split(b, b.shape[-1], axis=-1) + d1 = A11 + L_12 = A12 / d1 + d2 = A22 - L_12 * A12 + L_13 = A13 / d1 + L_23 = (A23 - L_13 * A12) / d2 + d3 = A33 - L_13 * A13 - L_23 * L_23 * d2 + y1 = b1 + y2 = b2 - L_12 * y1 + y3 = b3 - L_13 * y1 - L_23 * y2 + x3 = y3 / d3 + x2 = y2 / d2 - L_23 * x3 + x1 = y1 / d1 - L_12 * x2 - L_13 * x3 + return np.stack([x1, x2, x3], axis=-1).squeeze() + + +def weighted_downsample(X, confidence, scale=None, target_size=None): + if target_size is None: + target_size = (np.array(X.shape[:2]) / scale).round().astype(np.int) + if X.shape[1] > confidence.shape[1]: + X = resize(X, confidence.shape) + if X.ndim == 3: + confidence = confidence[..., np.newaxis] + numerator = resize(X * confidence, target_size) + denom = resize(confidence, target_size) + if X.ndim == 3: + denom = denom[..., np.newaxis] + return numerator / denom + + +def guided_upsample( + reference, + source, + kernel_size, + confidence=None, + eps_luma=1e-2, + eps_chroma=1e-2, + clip_output=True, +): + assert reference.shape[2] == 3 + + if np.any(np.array(source.shape) < np.array(reference.shape[:2])): + source = resize(source, reference.shape[:2]) + if confidence is None: + confidence = probability_to_confidence(source) + assert confidence.shape == source.shape + + reference_small = weighted_downsample(reference, confidence, kernel_size) + small_shape = reference_small.shape[:2] + source_small = weighted_downsample(source, confidence, target_size=small_shape) + + outer_reference = outer_product_images(reference, reference) + outer_reference = weighted_downsample(outer_reference, confidence, target_size=small_shape) + covar = outer_reference - outer_product_images(reference_small, reference_small) + var = weighted_downsample( + reference * source[..., np.newaxis], confidence, target_size=small_shape + ) + residual_small = var - reference_small * source_small[..., np.newaxis] + covar[..., 0] += eps_luma ** 2 + covar[..., [3, 5]] += eps_chroma ** 2 + + affine = solve_image_ldl3(covar, residual_small) + residual = source_small - (affine * reference_small).sum(axis=2) + affine = smooth_upsample(affine, reference.shape[:2]) + residual = smooth_upsample(residual, reference.shape[:2]) + output = (affine * reference).sum(axis=2) + residual + if clip_output: + output = output.clip(0, 1) + return output + + +if __name__ == "__main__": + img = cv2.imread("full-resolution-images/fig02/204315_img.jpg") + mask = cv2.imread("full-resolution-images/fig02/204315_inferred_mask.jpg") + img = img[..., ::-1] / 255 + mask = np.rot90(mask[..., 0], -1) / 255 + # we work with original resolution so kernel size multiplied by a factor of 4 + kernel_size = 64 * 4 + filtered = guided_upsample(img, mask, kernel_size) + cv2.imwrite("output.png", (255 * filtered).round().astype(np.uint8))