Files
bmad4ever-comfyui_quilting/make_seamless.py
T
Bruno Madeira 657dd120e1 cleaning ...
* patch_search.py renamed to synthesis_subroutines.py.
* fixed number of channels when using cv.floodFill with latents in new min cut implementation.
* make_seamless.py and make_seamless2.py use same "min_cut_patch" implementation.
* quilting.py allows to use v0 implementation, but otherwise, also shares the same as seamless nodes.
* check optional lib using importlib
2024-08-19 15:56:04 +01:00

152 lines
8.1 KiB
Python

from .synthesis_subroutines import get_generic_find_patch_method, get_4way_min_cut_patch
from .types import UiCoordData
from math import ceil
import numpy as np
def get_numb_of_blocks_to_fill_stripe(block_size, overlap, dim_length):
return int(ceil((dim_length - block_size) / (block_size - overlap)))
def make_seamless_horizontally(image, block_size, overlap, tolerance, rng: np.random.Generator,
version: int = 1, lookup_texture=None, uicd: UiCoordData | None = None):
"""
@param image: the image to make seamless; also be used to fetch the patches.
@param fnc: when evaluating potential patches the errors of different adjacency will be combined using this function
@param lookup_texture: if provided, the patches will be obtained from "lookup_texture" instead.
"""
lookup_texture = image if lookup_texture is None else lookup_texture
bmo = block_size - overlap
src_h, src_w = image.shape[:2]
n_h = get_numb_of_blocks_to_fill_stripe(block_size, overlap, src_h) #int(ceil((src_h - block_size) / bmo))
texture_map = np.zeros((src_h, src_w, image.shape[-1])).astype(image.dtype)
texture_map[:src_h, :src_w] = image
# roll texture map to allow overlapping between left and right blocks.
# this allows for big block sizes, and, consequently, for faster generations
left_blocks = np.roll(texture_map, block_size // 2 - overlap + block_size, axis=1)[:, :block_size]
right_blocks = np.roll(texture_map, -round(block_size / 2) + overlap, axis=1)[:, :block_size]
# center v seam at half block distance of the left corner
texture_map = np.roll(texture_map, block_size // 2, axis=1)
# get 1st patch
ref_block_left = left_blocks[:block_size, :block_size]
ref_block_right = right_blocks[:block_size, :block_size]
find_patch = get_generic_find_patch_method(version=version)
patch_block = find_patch(
ref_block_left, ref_block_right, None, None, lookup_texture, block_size, overlap, tolerance, rng) #, fnc)
min_cut_patch = get_4way_min_cut_patch(ref_block_left, ref_block_right, None, None,
patch_block, block_size, overlap)
texture_map[:block_size, :block_size] = min_cut_patch
if uicd is not None and uicd.add_to_job_data_slot_and_check_interrupt(1):
return None
for y in range(1, n_h):
blk_1y = y * bmo # block top corner y
blk_2y = blk_1y + block_size # block bottom corner y
# get adjacent blocks
ref_block_left = left_blocks[blk_1y:blk_2y, :block_size]
ref_block_right = right_blocks[blk_1y:blk_2y, :block_size]
ref_block_top = texture_map[(blk_1y - bmo):(blk_1y + overlap), :block_size]
patch_block = find_patch(ref_block_left, ref_block_right, ref_block_top, None,
lookup_texture, block_size, overlap, tolerance, rng)
min_cut_patch = get_4way_min_cut_patch(ref_block_left, ref_block_right, ref_block_top, None,
patch_block, block_size, overlap)
texture_map[blk_1y:blk_2y, :block_size] = min_cut_patch
if uicd is not None and uicd.add_to_job_data_slot_and_check_interrupt(1):
return None
# fill last block
ref_block_left = left_blocks[-block_size:, :block_size]
ref_block_right = right_blocks[-block_size:, :block_size]
ref_block_top = np.empty_like(ref_block_left) # only copy overlap
ref_block_top[-overlap:, :] = texture_map[-block_size:-block_size + overlap, :block_size]
patch_block = find_patch(ref_block_left, ref_block_right, ref_block_top, None,
lookup_texture, block_size, overlap, tolerance, rng)
min_cut_patch = get_4way_min_cut_patch(ref_block_left, ref_block_right, ref_block_top, None,
patch_block, block_size, overlap)
texture_map[-block_size:, :block_size] = min_cut_patch
if uicd is not None and uicd.add_to_job_data_slot_and_check_interrupt(1):
return None
return texture_map
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(
np.rot90(image, 1), block_size, overlap, tolerance, rng=rng, version=version, uicd=uicd,
lookup_texture=None if lookup_texture is None else np.rot90(lookup_texture))
return np.rot90(rotated_solution, -1).copy() if rotated_solution is not None else None
def make_seamless_both(image, block_size, overlap, tolerance, rng: np.random.Generator,
version: int = 1, lookup_texture=None, uicd: UiCoordData | None = None):
lookup_texture = image if lookup_texture is None else lookup_texture
# patch the texture in both directions. the last stripe's endpoints won't loop yet.
texture = make_seamless_vertically(image, block_size, overlap, tolerance, rng,
version=version, lookup_texture=lookup_texture, uicd=uicd)
if texture is not None:
texture = np.roll(texture, -block_size // 2, axis=0) # center future seam at stripes interception
texture = make_seamless_horizontally(texture, block_size, overlap, tolerance, rng,
version=version, lookup_texture=lookup_texture, uicd=uicd)
if texture is None:
return None
# center the area to patch 1st, this will make the rolls in the next step easier
texture = np.roll(texture, texture.shape[0] // 2, axis=0)
texture = np.roll(texture, (texture.shape[1] - block_size) // 2, axis=1)
return patch_horizontal_seam(texture, lookup_texture, block_size, overlap, tolerance, rng,
version=version, uicd=uicd)
def patch_horizontal_seam(texture_to_patch, lookup_texture, block_size, overlap, tolerance,
rng: np.random.Generator, version: int = 1, uicd: UiCoordData | None = None):
"""
Patches the center of the texture
"""
ys = (texture_to_patch.shape[0] - block_size) // 2
ye = ys + block_size
xs = (texture_to_patch.shape[1] - block_size) // 2
xe = xs + block_size
find_patch = get_generic_find_patch_method(version=version)
# PATCH H SEAM -> LEFT PATCH
adj_top_blk = np.roll(texture_to_patch, ye - overlap, axis=0)[-block_size:, xs - overlap:xe - overlap]
adj_btm_blk = np.roll(texture_to_patch, -ye + overlap, axis=0)[:block_size, xs - overlap:xe - overlap]
adj_lft_blk = np.roll(texture_to_patch, -xs, axis=1)[ys:ye, -block_size:]
patch = find_patch(adj_lft_blk, None, adj_top_blk, adj_btm_blk,
lookup_texture, block_size, overlap, tolerance, rng)
patch = get_4way_min_cut_patch(adj_lft_blk, None, adj_top_blk, adj_btm_blk,
patch, block_size, overlap)
texture_to_patch[ys:ye, xs - overlap:xe - overlap] = patch
if uicd is not None and uicd.add_to_job_data_slot_and_check_interrupt(1):
return None
# PATCH H SEAM -> RIGHT PATCH
adj_top_blk = np.roll(texture_to_patch, ye - overlap, axis=0)[-block_size:, xs + overlap:xe + overlap]
adj_btm_blk = np.roll(texture_to_patch, -ye + overlap, axis=0)[:block_size, xs + overlap:xe + overlap]
adj_lft_blk = np.roll(texture_to_patch, -xs - overlap * 2, axis=1)[ys:ye, -block_size:] # review this one
adj_rgt_blk = np.roll(texture_to_patch, -xs - block_size, axis=1)[ys:ye, :block_size]
patch = find_patch(adj_lft_blk, adj_rgt_blk, adj_top_blk, adj_btm_blk,
lookup_texture, block_size, overlap, tolerance, rng)
patch = get_4way_min_cut_patch(adj_lft_blk, adj_rgt_blk, adj_top_blk, adj_btm_blk,
patch, block_size, overlap)
texture_to_patch[ys:ye, xs + overlap:xe + overlap] = patch
if uicd is not None and uicd.add_to_job_data_slot_and_check_interrupt(1):
return None
return texture_to_patch