mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
155 lines
7.6 KiB
Python
155 lines
7.6 KiB
Python
import math
|
|
from typing import Union
|
|
|
|
import numpy as np
|
|
|
|
from invokeai.backend.tiles.utils import TBLR, Tile, paste
|
|
|
|
|
|
def calc_tiles_with_overlap(
|
|
image_height: int, image_width: int, tile_height: int, tile_width: int, overlap: int = 0
|
|
) -> list[Tile]:
|
|
"""Calculate the tile coordinates for a given image shape under a simple tiling scheme with overlaps.
|
|
|
|
Args:
|
|
image_height (int): The image height in px.
|
|
image_width (int): The image width in px.
|
|
tile_height (int): The tile height in px. All tiles will have this height.
|
|
tile_width (int): The tile width in px. All tiles will have this width.
|
|
overlap (int, optional): The target overlap between adjacent tiles. If the tiles do not evenly cover the image
|
|
shape, then the last row/column of tiles will overlap more than this. Defaults to 0.
|
|
|
|
Returns:
|
|
list[Tile]: A list of tiles that cover the image shape. Ordered from left-to-right, top-to-bottom.
|
|
"""
|
|
assert image_height >= tile_height
|
|
assert image_width >= tile_width
|
|
assert overlap < tile_height
|
|
assert overlap < tile_width
|
|
|
|
non_overlap_per_tile_height = tile_height - overlap
|
|
non_overlap_per_tile_width = tile_width - overlap
|
|
|
|
num_tiles_y = math.ceil((image_height - overlap) / non_overlap_per_tile_height)
|
|
num_tiles_x = math.ceil((image_width - overlap) / non_overlap_per_tile_width)
|
|
|
|
# tiles[y * num_tiles_x + x] is the tile for the y'th row, x'th column.
|
|
tiles: list[Tile] = []
|
|
|
|
# Calculate tile coordinates. (Ignore overlap values for now.)
|
|
for tile_idx_y in range(num_tiles_y):
|
|
for tile_idx_x in range(num_tiles_x):
|
|
tile = Tile(
|
|
coords=TBLR(
|
|
top=tile_idx_y * non_overlap_per_tile_height,
|
|
bottom=tile_idx_y * non_overlap_per_tile_height + tile_height,
|
|
left=tile_idx_x * non_overlap_per_tile_width,
|
|
right=tile_idx_x * non_overlap_per_tile_width + tile_width,
|
|
),
|
|
overlap=TBLR(top=0, bottom=0, left=0, right=0),
|
|
)
|
|
|
|
if tile.coords.bottom > image_height:
|
|
# If this tile would go off the bottom of the image, shift it so that it is aligned with the bottom
|
|
# of the image.
|
|
tile.coords.bottom = image_height
|
|
tile.coords.top = image_height - tile_height
|
|
|
|
if tile.coords.right > image_width:
|
|
# If this tile would go off the right edge of the image, shift it so that it is aligned with the
|
|
# right edge of the image.
|
|
tile.coords.right = image_width
|
|
tile.coords.left = image_width - tile_width
|
|
|
|
tiles.append(tile)
|
|
|
|
def get_tile_or_none(idx_y: int, idx_x: int) -> Union[Tile, None]:
|
|
if idx_y < 0 or idx_y > num_tiles_y or idx_x < 0 or idx_x > num_tiles_x:
|
|
return None
|
|
return tiles[idx_y * num_tiles_x + idx_x]
|
|
|
|
# Iterate over tiles again and calculate overlaps.
|
|
for tile_idx_y in range(num_tiles_y):
|
|
for tile_idx_x in range(num_tiles_x):
|
|
cur_tile = get_tile_or_none(tile_idx_y, tile_idx_x)
|
|
top_neighbor_tile = get_tile_or_none(tile_idx_y - 1, tile_idx_x)
|
|
left_neighbor_tile = get_tile_or_none(tile_idx_y, tile_idx_x - 1)
|
|
|
|
assert cur_tile is not None
|
|
|
|
# Update cur_tile top-overlap and corresponding top-neighbor bottom-overlap.
|
|
if top_neighbor_tile is not None:
|
|
cur_tile.overlap.top = max(0, top_neighbor_tile.coords.bottom - cur_tile.coords.top)
|
|
top_neighbor_tile.overlap.bottom = cur_tile.overlap.top
|
|
|
|
# Update cur_tile left-overlap and corresponding left-neighbor right-overlap.
|
|
if left_neighbor_tile is not None:
|
|
cur_tile.overlap.left = max(0, left_neighbor_tile.coords.right - cur_tile.coords.left)
|
|
left_neighbor_tile.overlap.right = cur_tile.overlap.left
|
|
|
|
return tiles
|
|
|
|
|
|
def merge_tiles_with_linear_blending(
|
|
dst_image: np.ndarray, tiles: list[Tile], tile_images: list[np.ndarray], blend_amount: int
|
|
):
|
|
"""Merge a set of image tiles into `dst_image` with linear blending between the tiles.
|
|
|
|
We expect every tile edge to either:
|
|
1) have an overlap of 0, because it is aligned with the image edge, or
|
|
2) have an overlap >= blend_amount.
|
|
If neither of these conditions are satisfied, we raise an exception.
|
|
|
|
The linear blending is centered at the halfway point of the overlap between adjacent tiles.
|
|
|
|
Args:
|
|
dst_image (np.ndarray): The destination image. Shape: (H, W, C).
|
|
tiles (list[Tile]): The list of tiles describing the locations of the respective `tile_images`.
|
|
tile_images (list[np.ndarray]): The tile images to merge into `dst_image`.
|
|
blend_amount (int): The amount of blending (in px) between adjacent overlapping tiles.
|
|
"""
|
|
# Sort tiles and images first by left x coordinate, then by top y coordinate. During tile processing, we want to
|
|
# iterate over tiles left-to-right, top-to-bottom.
|
|
tiles_and_images = list(zip(tiles, tile_images, strict=True))
|
|
tiles_and_images = sorted(tiles_and_images, key=lambda x: x[0].coords.left)
|
|
tiles_and_images = sorted(tiles_and_images, key=lambda x: x[0].coords.top)
|
|
|
|
# Prepare 1D linear gradients for blending.
|
|
gradient_left_x = np.linspace(start=0.0, stop=1.0, num=blend_amount)
|
|
gradient_top_y = np.linspace(start=0.0, stop=1.0, num=blend_amount)
|
|
# Convert shape: (blend_amount, ) -> (blend_amount, 1). The extra dimension enables the gradient to be applied
|
|
# to a 2D image via broadcasting. Note that no additional dimension is needed on gradient_left_x for
|
|
# broadcasting to work correctly.
|
|
gradient_top_y = np.expand_dims(gradient_top_y, axis=1)
|
|
|
|
for tile, tile_image in tiles_and_images:
|
|
# We expect tiles to be written left-to-right, top-to-bottom. We construct a mask that applies linear blending
|
|
# to the top and to the left of the current tile. The inverse linear blending is automatically applied to the
|
|
# bottom/right of the tiles that have already been pasted by the paste(...) operation.
|
|
tile_height, tile_width, _ = tile_image.shape
|
|
mask = np.ones(shape=(tile_height, tile_width), dtype=np.float64)
|
|
# Top blending:
|
|
if tile.overlap.top > 0:
|
|
assert tile.overlap.top >= blend_amount
|
|
# Center the blending gradient in the middle of the overlap.
|
|
blend_start_top = tile.overlap.top // 2 - blend_amount // 2
|
|
# The region above the blending region is masked completely.
|
|
mask[:blend_start_top, :] = 0.0
|
|
# Apply the blend gradient to the mask. Note that we use `*=` rather than `=` to achieve more natural
|
|
# behavior on the corners where vertical and horizontal blending gradients overlap.
|
|
mask[blend_start_top : blend_start_top + blend_amount, :] *= gradient_top_y
|
|
# For visual debugging:
|
|
# tile_image[blend_start_top : blend_start_top + blend_amount, :] = 0
|
|
|
|
# Left blending:
|
|
# (See comments under 'top blending' for an explanation of the logic.)
|
|
if tile.overlap.left > 0:
|
|
assert tile.overlap.left >= blend_amount
|
|
blend_start_left = tile.overlap.left // 2 - blend_amount // 2
|
|
mask[:, :blend_start_left] = 0.0
|
|
mask[:, blend_start_left : blend_start_left + blend_amount] *= gradient_left_x
|
|
# For visual debugging:
|
|
# tile_image[:, blend_start_left : blend_start_left + blend_amount] = 0
|
|
|
|
paste(dst_image=dst_image, src_image=tile_image, box=tile.coords, mask=mask)
|