nodes/bmad4ever-comfyui_quilting
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add patch_search.py, which implements patch_search using opencv matchTemplate.

Browse Source 
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 )
This commit is contained in:
Bruno Madeira
2024-07-09 01:36:46 +01:00
parent b6269a3c91
commit ec51afeba1
4 changed files with 259 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+30 -25
View File
@@ -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)
parallel_quilting.py
+111 -38
View File
@@ -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
patch_search.py
+116
View File
@@ -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]
requirements.txt
+2 -1
View File
@@ -1 +1,2 @@
joblib~=1.3.2
joblib~=1.3.2
opencv-python~=4.8.1.78