utils/image_utils.py

"""Image utilities for Indic Heritage Studio v2."""
from __future__ import annotations

from pathlib import Path
from typing import Tuple, Union

import numpy as np
from PIL import Image


def load_image(path: Union[str, Path]) -> Image.Image:
    """Load an image as RGB PIL."""
    img = Image.open(path)
    if img.mode != "RGB":
        img = img.convert("RGB")
    return img


def resize_to_sdxl(image: Image.Image, target: int = 1024) -> Image.Image:
    """Resize so the longer edge = target, preserve aspect ratio."""
    w, h = image.size
    if w >= h:
        new_w, new_h = target, int(h * target / w)
    else:
        new_w, new_h = int(w * target / h), target
    return image.resize((new_w, new_h), Image.LANCZOS)


def center_crop_to(image: Image.Image, size: Tuple[int, int]) -> Image.Image:
    """Center-crop an image to exact dimensions."""
    w, h = image.size
    tw, th = size
    left = (w - tw) // 2
    top = (h - th) // 2
    return image.crop((left, top, left + tw, top + th))


def image_to_numpy(image: Image.Image) -> np.ndarray:
    return np.array(image.convert("RGB"))


def numpy_to_image(arr: np.ndarray) -> Image.Image:
    return Image.fromarray(arr.astype(np.uint8))


def save_grid(images, path: Union[str, Path], cols: int = 4, cell_size: int = 256) -> None:
    """Save a list of images as a grid PNG."""
    from PIL import Image as PILImage
    n = len(images)
    rows = (n + cols - 1) // cols
    grid = PILImage.new("RGB", (cols * cell_size, rows * cell_size), color="white")
    for i, img in enumerate(images):
        r, c = i // cols, i % cols
        img_resized = img.resize((cell_size, cell_size), PILImage.LANCZOS)
        grid.paste(img_resized, (c * cell_size, r * cell_size))
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    grid.save(path)