diff --git a/make_seamless.py b/make_seamless.py index 96e4e63..1ab7c17 100644 --- a/make_seamless.py +++ b/make_seamless.py @@ -1,11 +1,8 @@ -from .patch_search import get_generic_find_patch_method, blur_patch_mask -from .misc.bse_type_aliases import num_pixels +from .patch_search import get_generic_find_patch_method, get_4way_min_cut_patch from .types import UiCoordData from math import ceil import numpy as np -inf = float('inf') - def get_numb_of_blocks_to_fill_stripe(block_size, overlap, dim_length): return int(ceil((dim_length - block_size) / (block_size - overlap))) @@ -84,108 +81,6 @@ def make_seamless_horizontally(image, block_size, overlap, tolerance, rng: np.ra return texture_map -def get_min_cut_patch_mask_horizontal(block1, block2, block_size: num_pixels, overlap: num_pixels): - """ - @param block1: block to the left, with the overlap on its right edge - @param block2: block to the right, with the overlap on its left edge - @return: ONLY the mask (not the patched overlap section) - """ - err = ((block1[:, -overlap:] - block2[:, :overlap]) ** 2).mean(2) - # maintain minIndex for 2nd row onwards and - min_index = [] - E = [list(err[0])] - for i in range(1, err.shape[0]): - # Get min values and args, -1 = left, 0 = middle, 1 = right - e = [inf] + E[-1] + [inf] - e = np.array([e[:-2], e[1:-1], e[2:]]) - # Get minIndex - min_arr = e.min(0) - min_arg = e.argmin(0) - 1 - min_index.append(min_arg) - # Set Eij = e_ij + min_ - Eij = err[i] + min_arr - E.append(list(Eij)) - - # Check the last element and backtrack to find path - path = [] - min_arg = np.argmin(E[-1]) - path.append(min_arg) - - # Backtrack to min path - for idx in min_index[::-1]: - min_arg = min_arg + idx[min_arg] - path.append(min_arg) - # Reverse to find full path - path = path[::-1] - mask = np.zeros((block_size, block_size, block1.shape[2]), dtype=block1.dtype) - for i in range(len(path)): - mask[i, :path[i] + 1] = 1 - return mask - - -def get_4way_min_cut_patch(ref_block_left, ref_block_right, ref_block_top, ref_block_bottom, - patch_block, block_size, overlap): - # (optional step) blur masks for a more seamless integration ( sometimes makes transition more noticeable, depends ) - masks_list = [] - - has_left = ref_block_left is None - has_right = ref_block_right is not None - has_top = ref_block_top is not None - has_bottom = ref_block_bottom is not None - - if has_left: - mask_left = get_min_cut_patch_mask_horizontal(ref_block_left, patch_block, block_size, overlap) - mask_left = blur_patch_mask(mask_left, block_size, overlap, has_left, has_right, has_top, has_bottom) - masks_list.append(mask_left) - - if has_right: - mask_right = get_min_cut_patch_mask_horizontal(np.fliplr(ref_block_right), np.fliplr(patch_block), block_size, - overlap) - mask_right = np.fliplr(mask_right) - mask_right = blur_patch_mask(mask_right, block_size, overlap, has_left, has_right, has_top, has_bottom) - masks_list.append(mask_right) - - if has_top: - # V , > counterclockwise rotation - mask_top = get_min_cut_patch_mask_horizontal(np.rot90(ref_block_top), np.rot90(patch_block), block_size, overlap) - mask_top = np.rot90(mask_top, 3) - mask_top = blur_patch_mask(mask_top, block_size, overlap, has_left, has_right, has_top, has_bottom) - masks_list.append(mask_top) - - if has_bottom: - mask_bottom = get_min_cut_patch_mask_horizontal(np.fliplr(np.rot90(ref_block_bottom)), - np.fliplr(np.rot90(patch_block)), block_size, overlap) - mask_bottom = np.rot90(np.fliplr(mask_bottom), 3) - mask_bottom = blur_patch_mask(mask_bottom, block_size, overlap, has_left, has_right, has_top, has_bottom) - masks_list.append(mask_bottom) - - # --- apply masks and return block --- - # compute auxiliary data - mask_s = sum(masks_list) - masks_max = np.maximum.reduce(masks_list) - mask_mos = np.divide(masks_max, mask_s, out=np.zeros_like(mask_s), where=mask_s != 0) - # note -> if blurred, masks weights are scaled with respect the max mask value. - # example: mask1: 0.2 mask2:0.5 -> max = .5 ; sum = .7 -> (.2*lb + .5*tb) * (max/sum) + patch * (1 - max) - - # place adjacent block sections - res_block = np.zeros_like(patch_block) - if has_left: - res_block += mask_left * np.roll(ref_block_left, overlap, 1) - if has_right: - res_block += mask_right * np.roll(ref_block_right, -overlap, 1) - if has_top: - res_block += mask_top * np.roll(ref_block_top, overlap, 0) - if has_bottom: - res_block += mask_bottom * np.roll(ref_block_bottom, -overlap, 0) - res_block *= mask_mos - - # place patch section - patch_weight = 1 - masks_max - res_block = res_block + patch_weight * patch_block - - return res_block - - def make_seamless_vertically(image, block_size, overlap, tolerance, rng: np.random.Generator, version: int = 1, lookup_texture=None, uicd: UiCoordData | None = None): rotated_solution = make_seamless_horizontally( diff --git a/patch_search.py b/patch_search.py index dec21f1..4987507 100644 --- a/patch_search.py +++ b/patch_search.py @@ -1,5 +1,4 @@ from functools import lru_cache -from math import ceil import numpy as np import cv2 as cv @@ -7,6 +6,7 @@ from .jena2020.generate import findPatchVertical, findPatchHorizontal, findPatch from .misc.bse_type_aliases import num_pixels epsilon = np.finfo(float).eps +inf = float('inf') # region get methods by version @@ -124,7 +124,7 @@ def find_patch_vx(ref_block_left, ref_block_right, ref_block_top, ref_block_bott @lru_cache(maxsize=4) def patch_blending_vignette(block_size: num_pixels, overlap: num_pixels, left: bool, right: bool, top: bool, bottom: bool): margin = 1 #ceil(overlap / 12) # must be small - power = 2 # controls drop-off + power = 2.5 # controls drop-off p = 6 # controls the shape def corner_distance(y, x): @@ -132,7 +132,7 @@ def patch_blending_vignette(block_size: num_pixels, overlap: num_pixels, left: b overlap - margin) return np.clip(distance, 0, 1) - mask = np.ones((block_size, block_size), dtype=float) + mask = np.ones((block_size, block_size), dtype=np.float32) i, j = np.meshgrid(np.arange(overlap), np.arange(overlap)) curve_top_left_corner = 1 - corner_distance(i, j) ** power @@ -149,24 +149,24 @@ def patch_blending_vignette(block_size: num_pixels, overlap: num_pixels, left: b if top and right: mask[:overlap, -overlap:] = np.flip(curve_top_left_corner, axis=1) elif top: - mask[:overlap, -overlap:] = curve_top_left_corner[:, -1].reshape(-1, 1) # Copy the last column flipped + mask[:overlap, -overlap:] = curve_top_left_corner[:, -1].reshape(-1, 1) elif right: - mask[:overlap, -overlap:] = np.flip(curve_top_left_corner[-1, :]).reshape(1, -1) # Copy the last row flipped + mask[:overlap, -overlap:] = np.flip(curve_top_left_corner[-1, :]).reshape(1, -1) # Bottom left corner if bottom and left: mask[-overlap:, :overlap] = np.flip(curve_top_left_corner, axis=0) elif bottom: - mask[-overlap:, :overlap] = np.flip(curve_top_left_corner[:, -1]).reshape(-1, 1) # Copy the last column flipped + mask[-overlap:, :overlap] = np.flip(curve_top_left_corner[:, -1]).reshape(-1, 1) elif left: - mask[-overlap:, :overlap] = curve_top_left_corner[-1, :] # Copy the last row flipped + mask[-overlap:, :overlap] = curve_top_left_corner[-1, :].reshape(1, -1) # Bottom right corner if bottom and right: mask[-overlap:, -overlap:] = np.flip(curve_top_left_corner) elif bottom: mask[-overlap:, -overlap:] = (np.flip(curve_top_left_corner[:, -1]) - .reshape(-1, 1)) # Copy the last column flipped + .reshape(-1, 1)) elif right: mask[-overlap:, -overlap:] = np.flip(curve_top_left_corner[-1, :]).reshape(1, -1) # Copy the last row flipped @@ -187,11 +187,11 @@ def patch_blending_vignette(block_size: num_pixels, overlap: num_pixels, left: b def blur_patch_mask(src_mask, block_size: num_pixels, overlap: num_pixels, left: bool, right: bool, top: bool, bottom: bool): + #print(f"src_mask type > {src_mask.dtype}") return src_mask # don't use it for now until further testing vignette = patch_blending_vignette(block_size, overlap, left, right, top, bottom) - # blurred = cv.GaussianBlur(src_mask, blur_ksize, blur_ksize) src_mask_uint = np.uint8(src_mask[:, :, 0] * 255) blurred = cv.distanceTransform(src_mask_uint, cv.DIST_L2, maskSize=0) blurred = cv.morphologyEx(blurred, cv.MORPH_ERODE, cv.getStructuringElement(cv.MORPH_ELLIPSE, (3, 3)), iterations=1) @@ -199,9 +199,153 @@ def blur_patch_mask(src_mask, block_size: num_pixels, overlap: num_pixels, left: blurred *= 255 / max(np.max(blurred), 1) blurred = np.float32(blurred / 255) - #blurred = np.stack((blurred,) * src_mask.shape[2], axis=-1) - #vignette = np.stack((vignette,) * src_mask.shape[2], axis=-1) - result = (vignette * blurred) + ((1 - vignette) * src_mask[:, :, 0]) + print(f"mask min max = {(np.min(result), np.max(result))}") + result = np.clip(result, 0, 1) # better safe than sorry result = np.stack((result, ) * src_mask.shape[2], axis=-1) return result + + +def get_min_cut_patch_mask_horizontal_jena2020(block1, block2, block_size: num_pixels, overlap: num_pixels): + """ + @param block1: block to the left, with the overlap on its right edge + @param block2: block to the right, with the overlap on its left edge + @return: ONLY the mask (not the patched overlap section) + """ + err = ((block1[:, -overlap:] - block2[:, :overlap]) ** 2).mean(2) + # maintain minIndex for 2nd row onwards and + min_index = [] + E = [list(err[0])] + for i in range(1, err.shape[0]): + # Get min values and args, -1 = left, 0 = middle, 1 = right + e = [inf] + E[-1] + [inf] + e = np.array([e[:-2], e[1:-1], e[2:]]) + # Get minIndex + min_arr = e.min(0) + min_arg = e.argmin(0) - 1 + min_index.append(min_arg) + # Set Eij = e_ij + min_ + Eij = err[i] + min_arr + E.append(list(Eij)) + + # Check the last element and backtrack to find path + path = [] + min_arg = np.argmin(E[-1]) + path.append(min_arg) + + # Backtrack to min path + for idx in min_index[::-1]: + min_arg = min_arg + idx[min_arg] + path.append(min_arg) + # Reverse to find full path + path = path[::-1] + mask = np.zeros((block_size, block_size, block1.shape[2]), dtype=block1.dtype) + for i in range(len(path)): + mask[i, :path[i] + 1] = 1 + return mask + + +def get_4way_min_cut_patch(ref_block_left, ref_block_right, ref_block_top, ref_block_bottom, + patch_block, block_size, overlap): + # (optional step) blur masks for a more seamless integration ( sometimes makes transition more noticeable, depends ) + masks_list = [] + + has_left = ref_block_left is not None + has_right = ref_block_right is not None + has_top = ref_block_top is not None + has_bottom = ref_block_bottom is not None + + if has_left: + mask_left = get_min_cut_patch_mask_horizontal(ref_block_left, patch_block, block_size, overlap) + mask_left = blur_patch_mask(mask_left, block_size, overlap, has_left, has_right, has_top, has_bottom) + masks_list.append(mask_left) + + if has_right: + mask_right = get_min_cut_patch_mask_horizontal(np.fliplr(ref_block_right), np.fliplr(patch_block), block_size, + overlap) + mask_right = np.fliplr(mask_right) + mask_right = blur_patch_mask(mask_right, block_size, overlap, has_left, has_right, has_top, has_bottom) + masks_list.append(mask_right) + + if has_top: + # V , > counterclockwise rotation + mask_top = get_min_cut_patch_mask_horizontal(np.rot90(ref_block_top), np.rot90(patch_block), block_size, overlap) + mask_top = np.rot90(mask_top, 3) + mask_top = blur_patch_mask(mask_top, block_size, overlap, has_left, has_right, has_top, has_bottom) + masks_list.append(mask_top) + + if has_bottom: + mask_bottom = get_min_cut_patch_mask_horizontal(np.fliplr(np.rot90(ref_block_bottom)), + np.fliplr(np.rot90(patch_block)), block_size, overlap) + mask_bottom = np.rot90(np.fliplr(mask_bottom), 3) + mask_bottom = blur_patch_mask(mask_bottom, block_size, overlap, has_left, has_right, has_top, has_bottom) + masks_list.append(mask_bottom) + + # --- apply masks and return block --- + # compute auxiliary data + mask_s = sum(masks_list) + masks_max = np.maximum.reduce(masks_list) + mask_mos = np.divide(masks_max, mask_s, out=np.zeros_like(mask_s), where=mask_s != 0) + # note -> if blurred, masks weights are scaled with respect the max mask value. + # example: mask1: 0.2 mask2:0.5 -> max = .5 ; sum = .7 -> (.2*lb + .5*tb) * (max/sum) + patch * (1 - max) + + # place adjacent block sections + res_block = np.zeros_like(patch_block) + if has_left: + res_block += mask_left * np.roll(ref_block_left, overlap, 1) + if has_right: + res_block += mask_right * np.roll(ref_block_right, -overlap, 1) + if has_top: + res_block += mask_top * np.roll(ref_block_top, overlap, 0) + if has_bottom: + res_block += mask_bottom * np.roll(ref_block_bottom, -overlap, 0) + res_block *= mask_mos + + # place patch section + patch_weight = 1 - masks_max + res_block = res_block + patch_weight * patch_block + return res_block + + +try: + import pyastar2d + + def get_min_cut_patch_mask_horizontal_astar(block1, block2, block_size: num_pixels, overlap: num_pixels): + """ + @param block1: block to the left, with the overlap on its right edge + @param block2: block to the right, with the overlap on its left edge + @return: ONLY the mask (not the patched overlap section) + """ + err = ((block1[:, -overlap:] - block2[:, :overlap]) ** 2).mean(2) + err *= block_size ** 3 + err += 1 + err *= block_size ** 3 # make the lowest value big enough for 1 to be negligible + err = np.pad(err, ((1, 1), (0, 0)), 'constant', constant_values=(1, 1)) + + start = (0, err.shape[1] // 2) + end = (err.shape[0] - 1, err.shape[1] // 2) + + path = pyastar2d.astar_path(err, start, end, allow_diagonal=True) + mask = np.ones((block_size, block_size, block1.shape[2]), dtype=block1.dtype) + shape_m2 = err.shape[0] - 2 + + start_index = 0 # find start index to avoid checking 0 < i every iteration + for idx, (i, j) in enumerate(path): + if 0 < i: + start_index = idx + break + + for i, j in path[start_index:]: # draw path + mask[i - 1, j + 1, :] = 0 + if i >= shape_m2: + break + + cv.floodFill(mask, None, (mask.shape[0]-1, mask.shape[1]-1), (0, ) * block1.shape[2]) + return mask + + + get_min_cut_patch_mask_horizontal = get_min_cut_patch_mask_horizontal_astar + print("comfyui_quilting: pyastar2d will be used to compute minimum cut.") +except Exception: + get_min_cut_patch_mask_horizontal = get_min_cut_patch_mask_horizontal_jena2020 + print("comfyui_quilting: jena2020 based solution will be used to compute minimum cut.")