Display the images in jupyter notebook

"""Smoke test for the tile infill"""

from pathlib import Path
from typing import Optional
from PIL import Image
from invokeai.backend.image_util.infill_methods.tile import infill_tile

images: list[tuple[str, Image.Image]] = []

for i in sorted(Path("./test_images/").glob("*.webp")):
    images.append((i.name, Image.open(i)))
    images.append((i.name, Image.open(i).transpose(Image.FLIP_LEFT_RIGHT)))
    images.append((i.name, Image.open(i).transpose(Image.FLIP_TOP_BOTTOM)))
    images.append((i.name, Image.open(i).resize((512, 512))))
    images.append((i.name, Image.open(i).resize((1234, 461))))

outputs: list[tuple[str, Image.Image, Image.Image, Optional[Image.Image]]] = []

for name, image in images:
    try:
        output = infill_tile(image, seed=0, tile_size=32)
        outputs.append((name, image, output.infilled, output.tile_image))
    except ValueError as e:
        print(f"Skipping image {name}: {e}")

# Display the images in jupyter notebook
import matplotlib.pyplot as plt
from PIL import ImageOps

fig, axes = plt.subplots(len(outputs), 3, figsize=(10, 3 * len(outputs)))
plt.subplots_adjust(hspace=0)

for i, (name, original, infilled, tile_image) in enumerate(outputs):
    # Add a border to each image, helps to see the edges
    size = original.size
    original = ImageOps.expand(original, border=5, fill="red")
    filled = ImageOps.expand(infilled, border=5, fill="red")
    if tile_image:
        tile_image = ImageOps.expand(tile_image, border=5, fill="red")

    axes[i, 0].imshow(original)
    axes[i, 0].axis("off")
    axes[i, 0].set_title(f"Original ({name} - {size})")

    if tile_image:
        axes[i, 1].imshow(tile_image)
        axes[i, 1].axis("off")
        axes[i, 1].set_title("Tile Image")
    else:
        axes[i, 1].axis("off")
        axes[i, 1].set_title("NO TILES GENERATED (NO TRANSPARENCY)")

    axes[i, 2].imshow(filled)
    axes[i, 2].axis("off")
    axes[i, 2].set_title("Filled")