mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
f0b1bb0327
The previous algorithm errored if the image wasn't divisible by the tile size. I've reimplemented it from scratch to mitigate this issue. The new algorithm is simpler. We create a pool of tiles, then use them to create an image composed completely of tiles. If there is any awkwardly sized space on the edge of the image, the tiles are cropped to fit. Finally, paste the original image over the tile image. I've added a jupyter notebook to do a smoke test of infilling methods, and 10 test images. The other infill algorithms can be easily tested with the notebook on the same images, though I didn't set that up yet. Tested and confirmed this gives results just as good as the earlier infill, though of course they aren't the same due to the change in the algorithm.
123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
def create_tile_pool(img_array: np.ndarray, tile_size: tuple[int, int]) -> list[np.ndarray]:
|
|
"""
|
|
Create a pool of tiles from non-transparent areas of the image by systematically walking through the image.
|
|
|
|
Args:
|
|
img_array: numpy array of the image.
|
|
tile_size: tuple (tile_width, tile_height) specifying the size of each tile.
|
|
|
|
Returns:
|
|
A list of numpy arrays, each representing a tile.
|
|
"""
|
|
tiles: list[np.ndarray] = []
|
|
rows, cols = img_array.shape[:2]
|
|
tile_width, tile_height = tile_size
|
|
|
|
for y in range(0, rows - tile_height + 1, tile_height):
|
|
for x in range(0, cols - tile_width + 1, tile_width):
|
|
tile = img_array[y : y + tile_height, x : x + tile_width]
|
|
# Check if the image has an alpha channel and the tile is completely opaque
|
|
if img_array.shape[2] == 4 and np.all(tile[:, :, 3] == 255):
|
|
tiles.append(tile)
|
|
elif img_array.shape[2] == 3: # If no alpha channel, append the tile
|
|
tiles.append(tile)
|
|
|
|
if not tiles:
|
|
raise ValueError(
|
|
"Not enough opaque pixels to generate any tiles. Use a smaller tile size or a different image."
|
|
)
|
|
|
|
return tiles
|
|
|
|
|
|
def create_filled_image(
|
|
img_array: np.ndarray, tile_pool: list[np.ndarray], tile_size: tuple[int, int], seed: int
|
|
) -> np.ndarray:
|
|
"""
|
|
Create an image of the same dimensions as the original, filled entirely with tiles from the pool.
|
|
|
|
Args:
|
|
img_array: numpy array of the original image.
|
|
tile_pool: A list of numpy arrays, each representing a tile.
|
|
tile_size: tuple (tile_width, tile_height) specifying the size of each tile.
|
|
|
|
Returns:
|
|
A numpy array representing the filled image.
|
|
"""
|
|
|
|
rows, cols, _ = img_array.shape
|
|
tile_width, tile_height = tile_size
|
|
|
|
# Prep an empty RGB image
|
|
filled_img_array = np.zeros((rows, cols, 3), dtype=img_array.dtype)
|
|
|
|
# Make the random tile selection reproducible
|
|
rng = np.random.default_rng(seed)
|
|
|
|
for y in range(0, rows, tile_height):
|
|
for x in range(0, cols, tile_width):
|
|
# Pick a random tile from the pool
|
|
tile = tile_pool[rng.integers(len(tile_pool))]
|
|
|
|
# Calculate the space available (may be less than tile size near the edges)
|
|
space_y = min(tile_height, rows - y)
|
|
space_x = min(tile_width, cols - x)
|
|
|
|
# Crop the tile if necessary to fit into the available space
|
|
cropped_tile = tile[:space_y, :space_x, :3]
|
|
|
|
# Fill the available space with the (possibly cropped) tile
|
|
filled_img_array[y : y + space_y, x : x + space_x, :3] = cropped_tile
|
|
|
|
return filled_img_array
|
|
|
|
|
|
@dataclass
|
|
class InfillTileOutput:
|
|
infilled: Image.Image
|
|
tile_image: Optional[Image.Image] = None
|
|
|
|
|
|
def infill_tile(image_to_infill: Image.Image, seed: int, tile_size: int) -> InfillTileOutput:
|
|
"""Infills an image with random tiles from the image itself.
|
|
|
|
If the image is not an RGBA image, it is returned untouched.
|
|
|
|
Args:
|
|
image: The image to infill.
|
|
tile_size: The size of the tiles to use for infilling.
|
|
|
|
Raises:
|
|
ValueError: If there are not enough opaque pixels to generate any tiles.
|
|
"""
|
|
|
|
if image_to_infill.mode != "RGBA":
|
|
return InfillTileOutput(infilled=image_to_infill)
|
|
|
|
# Internally, we want a tuple of (tile_width, tile_height). In the future, the tile size can be any rectangle.
|
|
_tile_size = (tile_size, tile_size)
|
|
np_image = np.array(image_to_infill, dtype=np.uint8)
|
|
|
|
# Create the pool of tiles that we will use to infill
|
|
tile_pool = create_tile_pool(np_image, _tile_size)
|
|
|
|
# Create an image from the tiles, same size as the original
|
|
tile_np_image = create_filled_image(np_image, tile_pool, _tile_size, seed)
|
|
|
|
# Paste the OG image over the tile image, effectively infilling the area
|
|
tile_image = Image.fromarray(tile_np_image, "RGB")
|
|
infilled = tile_image.copy()
|
|
infilled.paste(image_to_infill, (0, 0), image_to_infill.split()[-1])
|
|
|
|
# I think we want this to be "RGBA"?
|
|
infilled.convert("RGBA")
|
|
|
|
return InfillTileOutput(infilled=infilled, tile_image=tile_image)
|