inference/tiling.py

from pathlib import Path
from PIL import Image


DEFAULT_TILE_SIZE = 128
DEFAULT_OVERLAP = 32


def calculate_positions(length, tile_size, overlap):
    """
    Calculate tile starting positions along one dimension.

    The final tile always reaches the image boundary.
    """

    if tile_size <= 0:
        raise ValueError("tile_size must be positive.")

    if overlap < 0 or overlap >= tile_size:
        raise ValueError(
            "overlap must satisfy 0 <= overlap < tile_size."
        )

    stride = tile_size - overlap

    if length <= tile_size:
        return [0]

    positions = list(range(0, length - tile_size + 1, stride))

    # Ensure the final region is covered.
    last_position = length - tile_size

    if positions[-1] != last_position:
        positions.append(last_position)

    return positions


def create_tiles(
    image,
    tile_size=DEFAULT_TILE_SIZE,
    overlap=DEFAULT_OVERLAP
):
    """
    Split a PIL image into overlapping tiles.

    Returns:
        list of dictionaries containing:
            image
            x
            y
            width
            height
            index
    """

    if not isinstance(image, Image.Image):
        raise TypeError(
            "image must be a PIL.Image.Image"
        )

    image = image.convert("RGB")

    width, height = image.size

    x_positions = calculate_positions(
        width,
        tile_size,
        overlap
    )

    y_positions = calculate_positions(
        height,
        tile_size,
        overlap
    )

    tiles = []

    index = 0

    for y in y_positions:

        for x in x_positions:

            right = min(
                x + tile_size,
                width
            )

            bottom = min(
                y + tile_size,
                height
            )

            tile = image.crop(
                (x, y, right, bottom)
            )

            tiles.append(
                {
                    "image": tile,
                    "x": x,
                    "y": y,
                    "width": right - x,
                    "height": bottom - y,
                    "index": index
                }
            )

            index += 1

    return tiles


def save_tiles(
    tiles,
    output_dir,
    prefix="tile"
):
    """
    Save generated tiles to disk.
    """

    output_dir = Path(output_dir)
    output_dir.mkdir(
        parents=True,
        exist_ok=True
    )

    metadata = []

    for tile_info in tiles:

        image = tile_info["image"]

        filename = (
            f"{prefix}_"
            f"{tile_info['index']:03d}_"
            f"x{tile_info['x']}_"
            f"y{tile_info['y']}.png"
        )

        path = output_dir / filename

        image.save(path)

        metadata.append(
            {
                "index": tile_info["index"],
                "x": tile_info["x"],
                "y": tile_info["y"],
                "width": tile_info["width"],
                "height": tile_info["height"],
                "filename": filename
            }
        )

    return metadata