From ec51afeba17f8b5bdbf5ff2b20d0c76712fbf565 Mon Sep 17 00:00:00 2001 From: Bruno Madeira Date: Tue, 9 Jul 2024 01:36:46 +0100 Subject: [PATCH] Add patch_search.py, which implements patch_search using opencv matchTemplate. Updated requirements.txt to include opencv. Add patch_search option (named version) to existing nodes. Additional notes: - make_seamless.py should use the methods implemented in path_search.py (currently it is not, and its implementation has a bug). - consider importing cv only in functions, this way the prior solution can still be used w/o updating requirements... the speedup is so worth though... it might mislead potential users to keep using the 1st implementation... may should remove old solution and force opencv instead? ( to decide later, after seamless node implementation ) --- nodes.py | 55 ++++++++-------- parallel_quilting.py | 149 ++++++++++++++++++++++++++++++++----------- patch_search.py | 116 +++++++++++++++++++++++++++++++++ requirements.txt | 3 +- 4 files changed, 259 insertions(+), 64 deletions(-) create mode 100644 patch_search.py diff --git a/nodes.py b/nodes.py index 80af611..237d235 100644 --- a/nodes.py +++ b/nodes.py @@ -9,16 +9,16 @@ import torch # TODO add nodes where user defines output's height and width instead of scale -def quilt_single_src_no_parallelization(src, block_size, overlap, outH, outW, tolerance, rng: np.random.Generator, +def quilt_single_src_no_parallelization(src, block_size, overlap, outH, outW, tolerance, version, rng: np.random.Generator, jobs_shm_name, job_id): from .quilting.generate import generateTextureMap - return generateTextureMap(src, block_size, overlap, outH, outW, tolerance, rng, jobs_shm_name, job_id) + return generateTextureMap(src, block_size, overlap, outH, outW, tolerance, version, rng, jobs_shm_name, job_id) -def quilt_single_with_parallelization(src, block_size, overlap, outH, outW, tolerance, parallelization_lvl, rng, +def quilt_single_with_parallelization(src, block_size, overlap, outH, outW, tolerance, version, parallelization_lvl, rng, jobs_shm_name, job_id): from .parallel_quilting import generate_texture_parallel - return generate_texture_parallel(src, block_size, overlap, outH, outW, tolerance, parallelization_lvl, rng, + return generate_texture_parallel(src, block_size, overlap, outH, outW, tolerance, version, parallelization_lvl, rng, jobs_shm_name, job_id) @@ -29,13 +29,13 @@ QUILTING_SHARED_INPUT_TYPES = { # the percentage of pixels that overlap between each block_sized block "overlap": ("FLOAT", {"default": 1 / 6.0, "min": .1, "max": .9, "step": .01}), - # this is a percentage relative to the min error when searching for a patch. - # the ones within tolerance are potential candidates to be selected. + # this is a percentage relative to min error when searching for a patch. + # the ones that within tolerance are potential candidates to be selected. # tolerance equal to 1 means a tolerance of 2 times the min error. # my interpretation regarding its application is that # tolerance can help prevent too much sameness in the texture # due to some subset of patches being better at minimizing the error. - "tolerance": ("FLOAT", {"default": .1, "min": 0.01, "max": 2, "step": .01}), + "tolerance": ("FLOAT", {"default": .1, "min": 0, "max": 2, "step": .01}), # 0 -> no parallelization; uses the reference implementation # 1 -> 4 jobs used ( divides generation into 4 sections ) @@ -44,6 +44,8 @@ QUILTING_SHARED_INPUT_TYPES = { "parallelization_lvl": ("INT", {"default": 1, "min": 0, "max": 6, "step": 1}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + + "version": ("INT", {"default": 1, "min": 0, "max": 3}), } @@ -71,6 +73,7 @@ def waiting_loop(abort_loop_event: Event, pbar: utils.ProgressBar, total_steps, pbar.update_absolute(int(np.sum(procs_statuses[1:ntasks + 1])), total_steps) + def terminate_generation(finished_event, jobs_shared_memory, pbt: Thread): coord_jobs_array = np.ndarray((1,), dtype=np.dtype('uint32'), buffer=jobs_shared_memory.buf) interrupted = coord_jobs_array[0] @@ -128,7 +131,7 @@ class ImageQuilting: FUNCTION = "compute" CATEGORY = "Bmad/CV/Misc" - def compute(self, src, block_size, scale, overlap, tolerance, parallelization_lvl, seed): + def compute(self, src, block_size, scale, overlap, tolerance, parallelization_lvl, seed, version): h, w = src.shape[1:3] out_h, out_w = int(scale * h), int(scale * w) if overlap > 0: @@ -144,7 +147,7 @@ class ImageQuilting: setup_pbar(block_size, overlap, out_h, out_w, parallelization_lvl, src.shape[0]) if src.shape[0] > 1: # if image batch - texture_batch = self.batch_using_jobs(src, block_size, overlap, out_h, out_w, tolerance, + texture_batch = self.batch_using_jobs(src, block_size, overlap, out_h, out_w, tolerance, version, parallelization_lvl, rng, shm_name) terminate_generation(finish_event, shm_jobs, t) return (texture_batch,) @@ -152,11 +155,11 @@ class ImageQuilting: # if single image src = src.cpu().numpy().squeeze() if parallelization_lvl == 0: - texture = quilt_single_src_no_parallelization(src, block_size, overlap, out_h, out_w, tolerance, rng, - shm_name, 0) + texture = quilt_single_src_no_parallelization( + src, block_size, overlap, out_h, out_w, tolerance, version, rng, shm_name, 0) else: texture = quilt_single_with_parallelization( - src, block_size, overlap, out_h, out_w, tolerance, parallelization_lvl, rng, shm_name, 0) + src, block_size, overlap, out_h, out_w, tolerance, version, parallelization_lvl, rng, shm_name, 0) texture = torch.from_numpy(texture).unsqueeze(0) @@ -164,16 +167,16 @@ class ImageQuilting: return (texture,) @staticmethod - def batch_using_jobs(src, block_size, overlap, outH, outW, tolerance, parallelization_lvl, rng, jobs_shm_name): + def batch_using_jobs(src, block_size, overlap, outH, outW, tolerance, version, parallelization_lvl, rng, jobs_shm_name): from joblib import Parallel, delayed def unwrap_and_quilt(img_as_tensor, job_id): image = img_as_tensor.cpu().numpy() if parallelization_lvl == 0: - result = quilt_single_src_no_parallelization(image, block_size, overlap, outH, outW, tolerance, rng, + result = quilt_single_src_no_parallelization(image, block_size, overlap, outH, outW, tolerance, version, rng, jobs_shm_name, job_id) else: - result = quilt_single_with_parallelization(image, block_size, overlap, outH, outW, tolerance, 1, rng, + result = quilt_single_with_parallelization(image, block_size, overlap, outH, outW, tolerance, version, 1, rng, jobs_shm_name, job_id) return torch.from_numpy(result) @@ -201,7 +204,7 @@ class LatentQuilting: FUNCTION = "compute" CATEGORY = "Bmad/CV/Misc" - def compute(self, src, block_size, scale, overlap, tolerance, parallelization_lvl, seed): + def compute(self, src, block_size, scale, overlap, tolerance, parallelization_lvl, seed, version): src = src["samples"] h, w = src.shape[2:4] out_h, out_w = int(scale * h), int(scale * w) @@ -220,7 +223,7 @@ class LatentQuilting: setup_pbar(block_size, overlap, out_h, out_w, parallelization_lvl, src.shape[0]) if src.shape[0] > 1: # if multiple - latent_batch = self.batch_using_jobs(src, block_size, overlap, out_h, out_w, tolerance, + latent_batch = self.batch_using_jobs(src, block_size, overlap, out_h, out_w, tolerance, version, parallelization_lvl, rng, shm_name) terminate_generation(finish_event, shm_jobs, t) return ({"samples": latent_batch},) @@ -231,11 +234,11 @@ class LatentQuilting: if parallelization_lvl == 0: from .quilting.generate import generateTextureMap - texture = generateTextureMap(src, block_size, overlap, out_h, out_w, tolerance, rng, shm_name, 0) + texture = generateTextureMap(src, block_size, overlap, out_h, out_w, tolerance, version, rng, shm_name, 0) else: from .parallel_quilting import generate_texture_parallel - texture = generate_texture_parallel(src, block_size, overlap, out_h, out_w, tolerance, parallelization_lvl, rng, - shm_name, 0) + texture = generate_texture_parallel( + src, block_size, overlap, out_h, out_w, tolerance, version, parallelization_lvl, rng, shm_name, 0) terminate_generation(finish_event, shm_jobs, t) texture = np.moveaxis(texture, -1, 0) @@ -243,18 +246,20 @@ class LatentQuilting: return ({"samples": texture},) @staticmethod - def batch_using_jobs(src, block_size, overlap, outH, outW, tolerance, parallelization_lvl, rng, jobs_shm_name): + def batch_using_jobs(src, block_size, overlap, outH, outW, tolerance, version, parallelization_lvl, rng, jobs_shm_name): from joblib import Parallel, delayed def unwrap_and_quilt(latent, job_id): latent = latent.cpu().numpy().squeeze() latent = np.moveaxis(latent, 0, -1) if parallelization_lvl == 0: - result = quilt_single_src_no_parallelization(latent, block_size, overlap, outH, outW, tolerance, rng, - jobs_shm_name, job_id) + result = quilt_single_src_no_parallelization( + latent, block_size, overlap, outH, outW, tolerance, version, rng, + jobs_shm_name, job_id) else: - result = quilt_single_with_parallelization(latent, block_size, overlap, outH, outW, tolerance, 1, rng, - jobs_shm_name, job_id) + result = quilt_single_with_parallelization( + latent, block_size, overlap, outH, outW, tolerance, version, 1, rng, + jobs_shm_name, job_id) result = np.moveaxis(result, -1, 0) return torch.from_numpy(result) diff --git a/parallel_quilting.py b/parallel_quilting.py index 4433a5d..b35af05 100644 --- a/parallel_quilting.py +++ b/parallel_quilting.py @@ -1,8 +1,78 @@ from multiprocessing.shared_memory import SharedMemory + +import cv2 import numpy as np from math import ceil from .quilting.generate import findPatchVertical, findPatchHorizontal, findPatchBoth, \ getMinCutPatchHorizontal, getMinCutPatchVertical, getMinCutPatchBoth +from .patch_search import find_patch_v1, find_patch_v2, find_patch_v3 + + +# region get methods by version + +def get_find_patch_to_the_right_method(version: int): + match version: + case 0: + return findPatchHorizontal + case 1: + def v1_right(left_block, image, block_size, overlap, tolerance, rng): + return find_patch_v1(left_block, None, None, None, image, block_size, overlap, tolerance, rng) + + return v1_right + case 2: + def v2_right(left_block, image, block_size, overlap, tolerance, rng): + return find_patch_v2(left_block, None, None, None, image, block_size, overlap, tolerance, rng) + + return v2_right + case 3: + def v3_right(left_block, image, block_size, overlap, tolerance, rng): + return find_patch_v3(left_block, None, None, None, image, block_size, overlap, tolerance, rng) + + return v3_right + + +def get_find_patch_below_method(version: int): + match version: + case 0: + return findPatchVertical + case 1: + def v1_below(top_block, image, block_size, overlap, tolerance, rng): + return find_patch_v1(None, None, top_block, None, image, block_size, overlap, tolerance, rng) + + return v1_below + case 2: + def v2_below(top_block, image, block_size, overlap, tolerance, rng): + return find_patch_v2(None, None, top_block, None, image, block_size, overlap, tolerance, rng) + + return v2_below + case 3: + def v3_below(top_block, image, block_size, overlap, tolerance, rng): + return find_patch_v3(None, None, top_block, None, image, block_size, overlap, tolerance, rng) + + return v3_below + + +def get_find_patch_both_method(version: int): + match version: + case 0: + return findPatchBoth + case 1: + def v1_both(left_block, top_block, image, block_size, overlap, tolerance, rng): + return find_patch_v1(left_block, None, top_block, None, image, block_size, overlap, tolerance, rng) + + return v1_both + case 2: + def v2_both(left_block, top_block, image, block_size, overlap, tolerance, rng): + return find_patch_v2(left_block, None, top_block, None, image, block_size, overlap, tolerance, rng) + + return v2_both + case 3: + def v3_both(left_block, top_block, image, block_size, overlap, tolerance, rng): + return find_patch_v3(left_block, None, top_block, None, image, block_size, overlap, tolerance, rng) + + return v3_both + +# endregion # -- NOTE -- ___________________________________________________________________________________________________________ @@ -46,34 +116,38 @@ from .quilting.generate import findPatchVertical, findPatchHorizontal, findPatch # region methods extracted from the source of the referenced implementation for re-usability -def fill_column(image, initial_block, overlap, rows: int, tolerance, rng: np.random.Generator): +def fill_column(image, initial_block, overlap, rows: int, tolerance, version, rng: np.random.Generator): + find_patch_below = get_find_patch_below_method(version) block_size = initial_block.shape[0] texture_map = np.zeros( - ((block_size + rows * (block_size - overlap)), block_size, image.shape[2])) + ((block_size + rows * (block_size - overlap)), block_size, image.shape[2])).astype(image.dtype) texture_map[:block_size, :block_size, :] = initial_block for i, blkIdx in enumerate(range((block_size - overlap), texture_map.shape[0] - overlap, (block_size - overlap))): ref_block = texture_map[(blkIdx - block_size + overlap):(blkIdx + overlap), :block_size] - patch_block = findPatchVertical(ref_block, image, block_size, overlap, tolerance, rng) + patch_block = find_patch_below(ref_block, image, block_size, overlap, tolerance, rng) min_cut_patch = getMinCutPatchVertical(ref_block, patch_block, block_size, overlap) texture_map[blkIdx:(blkIdx + block_size), :block_size] = min_cut_patch return texture_map -def fill_row(image, initial_block, overlap, columns: int, tolerance, rng: np.random.Generator): +def fill_row(image, initial_block, overlap, columns: int, tolerance, version: int, rng: np.random.Generator): + find_patch_to_the_right = get_find_patch_to_the_right_method(version) + #left_block, image, block_size, overlap, tolerance, rng block_size = initial_block.shape[0] texture_map = np.zeros( - (block_size, (block_size + columns * (block_size - overlap)), image.shape[2])) + (block_size, (block_size + columns * (block_size - overlap)), image.shape[2])).astype(image.dtype) texture_map[:block_size, :block_size, :] = initial_block for i, blkIdx in enumerate(range((block_size - overlap), texture_map.shape[1] - overlap, (block_size - overlap))): ref_block = texture_map[:block_size, (blkIdx - block_size + overlap):(blkIdx + overlap)] - patch_block = findPatchHorizontal(ref_block, image, block_size, overlap, tolerance, rng) + patch_block = find_patch_to_the_right(ref_block, image, block_size, overlap, tolerance, rng) min_cut_patch = getMinCutPatchHorizontal(ref_block, patch_block, block_size, overlap) texture_map[:block_size, blkIdx:(blkIdx + block_size)] = min_cut_patch return texture_map -def fill_quad(rows: int, columns: int, block_size, overlap, texture_map, image, tolerance, +def fill_quad(rows: int, columns: int, block_size, overlap, texture_map, image, tolerance, version, rng: np.random.Generator, jobs_shm_name, job_id): + find_patch_both = get_find_patch_both_method(version) # here job id is more granular, not the same as generate_texture_parallel, but rather within the sub job shm_jobs = SharedMemory(name=jobs_shm_name) coord_jobs_array = np.ndarray((2 + job_id,), dtype=np.dtype('uint32'), buffer=shm_jobs.buf) @@ -89,7 +163,7 @@ def fill_quad(rows: int, columns: int, block_size, overlap, texture_map, image, (blk_index_i - block_size + overlap):(blk_index_i + overlap), blk_index_j:(blk_index_j + block_size)] - patch_block = findPatchBoth(ref_block_left, ref_block_top, image, block_size, overlap, tolerance, rng) + patch_block = find_patch_both(ref_block_left, ref_block_top, image, block_size, overlap, tolerance, rng) min_cut_patch = getMinCutPatchBoth(ref_block_left, ref_block_top, patch_block, block_size, overlap) texture_map[blk_index_i:(blk_index_i + block_size), blk_index_j:(blk_index_j + block_size)] = min_cut_patch @@ -103,8 +177,7 @@ def fill_quad(rows: int, columns: int, block_size, overlap, texture_map, image, # endregion - -def generate_texture_parallel(image, block_size, overlap, outH, outW, tolerance, nps, +def generate_texture_parallel(image, block_size, overlap, outH, outW, tolerance, version: int, nps, rng: np.random.Generator, jobs_shm_name, job_id): """ @param jobs_shm_name: shared memory name to a one dimensional array that stores the number of "blocks" @@ -143,10 +216,10 @@ def generate_texture_parallel(image, block_size, overlap, outH, outW, tolerance, # generate 2 vertical strips and 2 horizontal strips that will split the generated canvas in half # the center, where the stripes connect, shares the same tile args = [ - (image, start_block, overlap, cols_per_quad, tolerance, rng), - (hi_image, hi_start_block, overlap, cols_per_quad, tolerance, rng), - (image, start_block, overlap, rows_per_quad, tolerance, rng), - (vi_image, vi_start_block, overlap, rows_per_quad, tolerance, rng) + (image, start_block, overlap, cols_per_quad, tolerance, version, rng), + (hi_image, hi_start_block, overlap, cols_per_quad, tolerance, version, rng), + (image, start_block, overlap, rows_per_quad, tolerance, version, rng), + (vi_image, vi_start_block, overlap, rows_per_quad, tolerance, version, rng) ] funcs = [fill_row, fill_row, fill_column, fill_column] stripes = Parallel(n_jobs=4, backend="loky", timeout=None)( @@ -156,17 +229,17 @@ def generate_texture_parallel(image, block_size, overlap, outH, outW, tolerance, # generate the 4 sections (quadrants) args = [ - (vis, his, hi_image, rows_per_quad, cols_per_quad, overlap, tolerance, nps, rng, jobs_shm_name, job_id), - (vis, hs, vi_image, rows_per_quad, cols_per_quad, overlap, tolerance, nps, rng, jobs_shm_name, job_id), - (vs, hs, image, rows_per_quad, cols_per_quad, overlap, tolerance, nps, rng, jobs_shm_name, job_id), - (vs, his, hi_image, rows_per_quad, cols_per_quad, overlap, tolerance, nps, rng, jobs_shm_name, job_id) + (vis, his, hi_image, rows_per_quad, cols_per_quad, overlap, tolerance, version, nps, rng, jobs_shm_name, job_id), + (vis, hs, vi_image, rows_per_quad, cols_per_quad, overlap, tolerance, version, nps, rng, jobs_shm_name, job_id), + (vs, hs, image, rows_per_quad, cols_per_quad, overlap, tolerance, version, nps, rng, jobs_shm_name, job_id), + (vs, his, hi_image, rows_per_quad, cols_per_quad, overlap, tolerance, version, nps, rng, jobs_shm_name, job_id) ] funcs = [quad1, quad2, quad3, quad4] quads = Parallel(n_jobs=4, backend="loky", timeout=None)( delayed(funcs[i])(*args[i]) for i in range(4)) q1, q2, q3, q4 = quads - texture = np.zeros((q1.shape[0] * 2 - block_size, q1.shape[1] * 2 - block_size, image.shape[2])) + texture = np.zeros((q1.shape[0] * 2 - block_size, q1.shape[1] * 2 - block_size, image.shape[2])).astype(image.dtype) bmo = block_size - overlap texture[:q1.shape[0] - bmo, :q1.shape[1] - bmo] = q1[:q1.shape[0] - bmo, :q1.shape[1] - bmo] texture[:q1.shape[0] - bmo, q1.shape[1] - bmo:] = q2[:q1.shape[0] - bmo, overlap:] @@ -176,7 +249,7 @@ def generate_texture_parallel(image, block_size, overlap, outH, outW, tolerance, return texture[:outH, :outW] -def quad1(vis, his, hi_image, rows: int, columns: int, overlap, tolerance, p_strips, rng, jobs_shm_name, job_id): +def quad1(vis, his, hi_image, rows: int, columns: int, overlap, tolerance, version, p_strips, rng, jobs_shm_name, job_id): """ :param his: horizontal inverted stripe :param vis: vertical inverted stripe @@ -191,15 +264,15 @@ def quad1(vis, his, hi_image, rows: int, columns: int, overlap, tolerance, p_str shm_text = SharedMemory(create=True, size=size) texture = np.ndarray((vis.shape[0], his.shape[1], hi_image.shape[2]), dtype=hi_image.dtype, buffer=shm_text.buf) else: - texture = np.zeros((vis.shape[0], his.shape[1], hi_image.shape[2])) + texture = np.zeros((vis.shape[0], his.shape[1], hi_image.shape[2])).astype(hi_image.dtype) texture[:vi_hi_s.shape[0], :vi_hi_s.shape[1]] = vi_hi_s[:, :] texture[vi_hi_s.shape[0]:hi_vi_s.shape[0], :hi_vi_s.shape[1]] = hi_vi_s[vi_hi_s.shape[0]:, :] if p_strips > 1: - fill_quad_ps(rows, columns, vi_hi_s.shape[0], overlap, shm_text.name, vhi_image, tolerance, p_strips, rng, + fill_quad_ps(rows, columns, vi_hi_s.shape[0], overlap, shm_text.name, vhi_image, tolerance, version, p_strips, rng, jobs_shm_name, job_id + p_strips * 0) else: - texture = fill_quad(rows, columns, vi_hi_s.shape[0], overlap, texture, vhi_image, tolerance, rng, jobs_shm_name, + texture = fill_quad(rows, columns, vi_hi_s.shape[0], overlap, texture, vhi_image, tolerance, version, rng, jobs_shm_name, job_id + 0) texture = np.ascontiguousarray(np.flip(texture, axis=(0, 1))) @@ -209,7 +282,7 @@ def quad1(vis, his, hi_image, rows: int, columns: int, overlap, tolerance, p_str return texture -def quad2(vis, hs, vi_image, rows: int, columns: int, overlap, tolerance, p_strips, rng, jobs_shm_name, job_id): +def quad2(vis, hs, vi_image, rows: int, columns: int, overlap, tolerance, version, p_strips, rng, jobs_shm_name, job_id): shm_text = None vi_hs = np.ascontiguousarray(np.flipud(hs)) @@ -218,15 +291,15 @@ def quad2(vis, hs, vi_image, rows: int, columns: int, overlap, tolerance, p_stri shm_text = SharedMemory(create=True, size=size) texture = np.ndarray((vis.shape[0], hs.shape[1], vi_image.shape[2]), dtype=vi_image.dtype, buffer=shm_text.buf) else: - texture = np.zeros((vis.shape[0], hs.shape[1], vi_image.shape[2])) + texture = np.zeros((vis.shape[0], hs.shape[1], vi_image.shape[2])).astype(vi_image.dtype) texture[:hs.shape[0], :hs.shape[1]] = vi_hs[:, :] texture[hs.shape[0]:vis.shape[0], :vis.shape[1]] = vis[hs.shape[0]:, :] if p_strips > 1: - fill_quad_ps(rows, columns, hs.shape[0], overlap, shm_text.name, vi_image, tolerance, p_strips, rng, + fill_quad_ps(rows, columns, hs.shape[0], overlap, shm_text.name, vi_image, tolerance, version, p_strips, rng, jobs_shm_name, job_id + p_strips * 1) else: - texture = fill_quad(rows, columns, hs.shape[0], overlap, texture, vi_image, tolerance, rng, jobs_shm_name, + texture = fill_quad(rows, columns, hs.shape[0], overlap, texture, vi_image, tolerance, version, rng, jobs_shm_name, job_id + 1) texture = np.ascontiguousarray(np.flipud(texture)) @@ -236,7 +309,7 @@ def quad2(vis, hs, vi_image, rows: int, columns: int, overlap, tolerance, p_stri return texture -def quad4(vs, his, hi_image, rows: int, columns: int, overlap, tolerance, p_strips, rng, jobs_shm_name, job_id): +def quad4(vs, his, hi_image, rows: int, columns: int, overlap, tolerance, version, p_strips, rng, jobs_shm_name, job_id): shm_text = None hi_vs = np.ascontiguousarray(np.fliplr(vs)) @@ -245,15 +318,15 @@ def quad4(vs, his, hi_image, rows: int, columns: int, overlap, tolerance, p_stri shm_text = SharedMemory(create=True, size=size) texture = np.ndarray((vs.shape[0], his.shape[1], hi_image.shape[2]), dtype=hi_image.dtype, buffer=shm_text.buf) else: - texture = np.zeros((vs.shape[0], his.shape[1], hi_image.shape[2])) + texture = np.zeros((vs.shape[0], his.shape[1], hi_image.shape[2])).astype(hi_image.dtype) texture[:his.shape[0], :his.shape[1]] = his[:, :] texture[his.shape[0]:vs.shape[0], :vs.shape[1]] = hi_vs[his.shape[0]:, :] if p_strips > 1: - fill_quad_ps(rows, columns, his.shape[0], overlap, shm_text.name, hi_image, tolerance, p_strips, rng, + fill_quad_ps(rows, columns, his.shape[0], overlap, shm_text.name, hi_image, tolerance, version, p_strips, rng, jobs_shm_name, job_id + p_strips * 2) else: - texture = fill_quad(rows, columns, his.shape[0], overlap, texture, hi_image, tolerance, rng, jobs_shm_name, + texture = fill_quad(rows, columns, his.shape[0], overlap, texture, hi_image, tolerance, version, rng, jobs_shm_name, job_id + 2) texture = np.ascontiguousarray(np.fliplr(texture)) @@ -263,7 +336,7 @@ def quad4(vs, his, hi_image, rows: int, columns: int, overlap, tolerance, p_stri return texture -def quad3(vs, hs, image, rows: int, columns: int, overlap, tolerance, p_strips, rng, jobs_shm_name, job_id): +def quad3(vs, hs, image, rows: int, columns: int, overlap, tolerance, version, p_strips, rng, jobs_shm_name, job_id): shm_text = None if p_strips > 1: @@ -271,15 +344,15 @@ def quad3(vs, hs, image, rows: int, columns: int, overlap, tolerance, p_strips, shm_text = SharedMemory(create=True, size=size) texture = np.ndarray((vs.shape[0], hs.shape[1], image.shape[2]), dtype=image.dtype, buffer=shm_text.buf) else: - texture = np.zeros((vs.shape[0], hs.shape[1], image.shape[2])) + texture = np.zeros((vs.shape[0], hs.shape[1], image.shape[2])).astype(image.dtype) texture[:hs.shape[0], :hs.shape[1]] = hs[:, :] texture[hs.shape[0]:vs.shape[0], :vs.shape[1]] = vs[hs.shape[0]:, :] if p_strips > 1: - fill_quad_ps(rows, columns, vs.shape[1], overlap, shm_text.name, image, tolerance, p_strips, rng, + fill_quad_ps(rows, columns, vs.shape[1], overlap, shm_text.name, image, tolerance, version, p_strips, rng, jobs_shm_name, job_id + p_strips * 3) else: - return fill_quad(rows, columns, vs.shape[1], overlap, texture, image, tolerance, rng, jobs_shm_name, job_id + 3) + return fill_quad(rows, columns, vs.shape[1], overlap, texture, image, tolerance, version, rng, jobs_shm_name, job_id + 3) texture = texture.copy() if p_strips > 1: @@ -288,7 +361,7 @@ def quad3(vs, hs, image, rows: int, columns: int, overlap, tolerance, p_strips, return texture -def fill_quad_ps(rows, columns, block_size, overlap, texture_shared_mem_name, image, tolerance, total_procs, rng, +def fill_quad_ps(rows, columns, block_size, overlap, version, texture_shared_mem_name, image, tolerance, total_procs, rng, jobs_shm_name, job_id): from joblib import Parallel, delayed @@ -307,6 +380,7 @@ def fill_quad_ps(rows, columns, block_size, overlap, texture_shared_mem_name, im np_coord[2 * ip + 1] = 1 def fill_rows(pid, coord_shared_list_name, texture_shm_name, sub_job_id): + find_patch_both = get_find_patch_both_method(version) shm_jobs = SharedMemory(name=jobs_shm_name) coord_jobs_array = np.ndarray((2 + sub_job_id,), dtype=np.dtype('uint32'), buffer=shm_jobs.buf) @@ -338,7 +412,7 @@ def fill_quad_ps(rows, columns, block_size, overlap, texture_shared_mem_name, im (blk_index_i - block_size + overlap):(blk_index_i + overlap), blk_index_j:(blk_index_j + block_size)] - patch_block = findPatchBoth(ref_block_left, ref_block_top, image, block_size, overlap, tolerance, rng) + patch_block = find_patch_both(ref_block_left, ref_block_top, image, block_size, overlap, tolerance, rng) min_cut_patch = getMinCutPatchBoth(ref_block_left, ref_block_top, patch_block, block_size, overlap) texture[blk_index_i:(blk_index_i + block_size), blk_index_j:(blk_index_j + block_size)] = min_cut_patch @@ -352,4 +426,3 @@ def fill_quad_ps(rows, columns, block_size, overlap, texture_shared_mem_name, im shm_coord.close() shm_coord.unlink() - return diff --git a/patch_search.py b/patch_search.py new file mode 100644 index 0000000..5e03795 --- /dev/null +++ b/patch_search.py @@ -0,0 +1,116 @@ +import numpy as np +import cv2 as cv + +epsilon = np.finfo(float).eps + + +def find_patch_v1(ref_block_left, ref_block_right, ref_block_top, ref_block_bottom, + texture, block_size, overlap, tolerance, + rng: np.random.Generator + ): + """ + Re-implementation of the version 1.0 solution using matchTemplate to improve performance. + Uses the total instead of the mean for the errors matrix; other than that should be exactly the same. + Does not output the same as version 1.0. + """ + blks_sqdiffs = [] + if ref_block_left is not None: + blks_sqdiffs.append(cv.matchTemplate( + image=texture[:, :-block_size + overlap], + templ=ref_block_left[:, -overlap:], method=cv.TM_SQDIFF)) + if ref_block_right is not None: + blks_sqdiffs.append(cv.matchTemplate( + image=np.roll(texture, -block_size + overlap, axis=1)[:, :-block_size + overlap], + templ=ref_block_right[:, :overlap], method=cv.TM_SQDIFF)) + if ref_block_top is not None: + blks_sqdiffs.append(cv.matchTemplate( + image=texture[:-block_size + overlap, :], + templ=ref_block_top[-overlap:, :], method=cv.TM_SQDIFF)) + if ref_block_bottom is not None: + blks_sqdiffs.append(cv.matchTemplate( + image=np.roll(texture, -block_size + overlap, axis=0)[:-block_size + overlap, :], + templ=ref_block_bottom[:overlap, :], method=cv.TM_SQDIFF)) + + err_mat = np.add.reduce(blks_sqdiffs) + min_val = np.min(err_mat[err_mat > 0 if tolerance > 0 else True]) # ignore zeroes to enforce tolerance usage + y, x = np.nonzero(err_mat <= (1.0 + tolerance) * min_val) + c = rng.integers(len(y)) + y, x = y[c], x[c] + return texture[y:y + block_size, x:x + block_size] + + +def find_patch_v2(ref_block_left, ref_block_right, ref_block_top, ref_block_bottom, + texture, block_size, overlap, tolerance, + rng: np.random.Generator + ): + """ + Same as find_patch_v1 but chooses maximum error instead of the sum of errors, + when patching with multiple adjacent blocks. + """ + blks_sqdiffs = [] + if ref_block_left is not None: + blks_sqdiffs.append(cv.matchTemplate( + image=texture[:, :-block_size + overlap], + templ=ref_block_left[:, -overlap:], method=cv.TM_SQDIFF)) + if ref_block_right is not None: + blks_sqdiffs.append(cv.matchTemplate( + image=np.roll(texture, -block_size + overlap, axis=1)[:, :-block_size + overlap], + templ=ref_block_right[:, :overlap], method=cv.TM_SQDIFF)) + if ref_block_top is not None: + blks_sqdiffs.append(cv.matchTemplate( + image=texture[:-block_size + overlap, :], + templ=ref_block_top[-overlap:, :], method=cv.TM_SQDIFF)) + if ref_block_bottom is not None: + blks_sqdiffs.append(cv.matchTemplate( + image=np.roll(texture, -block_size + overlap, axis=0)[:-block_size + overlap, :], + templ=ref_block_bottom[:overlap, :], method=cv.TM_SQDIFF)) + + err_mat = np.maximum.reduce(blks_sqdiffs) + min_val = np.min(err_mat[err_mat > 0 if tolerance > 0 else True]) # ignore zeroes to enforce tolerance usage + y, x = np.nonzero(err_mat <= (1.0 + tolerance) * min_val) + c = rng.integers(len(y)) + y, x = y[c], x[c] + return texture[y:y + block_size, x:x + block_size] + + +def find_patch_v3(ref_block_left, ref_block_right, ref_block_top, ref_block_bottom, + texture, block_size, overlap, tolerance, + rng: np.random.Generator + ): + """ + This version makes use of TM_CCOEFF in matchTemplate instead of TM_SQDIFF. + """ + blks_ccs = [] + t_max = 0 + if ref_block_left is not None: + blk_overlap = ref_block_left[:, -overlap:] + blks_ccs.append(cv.matchTemplate( + image=texture[:, :-block_size + overlap], + templ=blk_overlap, method=cv.TM_CCOEFF)) + t_max = cv.matchTemplate(blk_overlap, blk_overlap, method=cv.TM_CCOEFF)[0] + if ref_block_right is not None: + blk_overlap = ref_block_right[:, :overlap] + blks_ccs.append(cv.matchTemplate( + image=np.roll(texture, -block_size + overlap, axis=1)[:, :-block_size + overlap], + templ=blk_overlap, method=cv.TM_CCOEFF)) + t_max += cv.matchTemplate(blk_overlap, blk_overlap, method=cv.TM_CCOEFF)[0] + if ref_block_top is not None: + blk_overlap = ref_block_top[-overlap:, :] + blks_ccs.append(cv.matchTemplate( + image=texture[:-block_size + overlap, :], + templ=blk_overlap, method=cv.TM_CCOEFF)) + t_max += cv.matchTemplate(blk_overlap, blk_overlap, method=cv.TM_CCOEFF)[0] + if ref_block_bottom is not None: + blk_overlap = ref_block_bottom[:overlap, :] + blks_ccs.append(cv.matchTemplate( + image=np.roll(texture, -block_size + overlap, axis=0)[:-block_size + overlap, :], + templ=blk_overlap, method=cv.TM_CCOEFF)) + t_max += cv.matchTemplate(blk_overlap, blk_overlap, method=cv.TM_CCOEFF)[0] + + err_mat = t_max - np.add.reduce(blks_ccs) + print(f"min val = {np.min(err_mat)}") + min_val = np.min(err_mat[err_mat > 0 if tolerance > 0 else True]) # ignore zeroes to enforce tolerance usage + y, x = np.nonzero(err_mat <= (1.0 + tolerance) * min_val) + c = rng.integers(len(y)) + y, x = y[c], x[c] + return texture[y:y + block_size, x:x + block_size] diff --git a/requirements.txt b/requirements.txt index f4059c9..f684cd1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ -joblib~=1.3.2 \ No newline at end of file +joblib~=1.3.2 +opencv-python~=4.8.1.78 \ No newline at end of file