data/preprocessing.py

"""Image preprocessing and visualization utilities."""

from __future__ import annotations

import hashlib
import io
import base64
from pathlib import Path
from typing import Any, Iterable

from PIL import Image, ImageDraw, ImageFont, ImageOps

from data.schemas import LABEL_DISPLAY_NAMES, bbox_to_pixels

LABEL_STYLE = {
    "dust": ((245, 158, 11), 2),
    "dirt": ((217, 119, 6), 2),
    "scratch": ((220, 38, 38), 3),
    "long_hair": ((124, 58, 237), 2),
    "short_hair": ((8, 145, 178), 2),
    "emulsion_damage": ((226, 232, 240), 3),
    "chemical_stain": ((22, 163, 74), 3),
    "light_leak": ((244, 114, 182), 3),
}
DEFAULT_STYLE = ((255, 255, 255), 2)


def load_image(image: str | Path | Image.Image) -> Image.Image:
    """Load an image-like value and return RGB PIL Image."""
    if isinstance(image, Image.Image):
        pil = image
    else:
        pil = Image.open(image)
    pil = ImageOps.exif_transpose(pil)
    if pil.mode == "RGBA":
        background = Image.new("RGB", pil.size, (24, 22, 20))
        background.paste(pil, mask=pil.getchannel("A"))
        return background
    if pil.mode != "RGB":
        return pil.convert("RGB")
    return pil.copy()


def image_to_png_bytes(image: Image.Image) -> bytes:
    buf = io.BytesIO()
    load_image(image).save(buf, format="PNG", optimize=True)
    return buf.getvalue()


def image_to_data_uri(
    image: Image.Image,
    *,
    max_side: int = 1800,
    image_format: str = "JPEG",
    quality: int = 92,
) -> str:
    """Return a browser-openable image data URI for review previews."""
    pil = resize_for_preview(load_image(image), max_side=max_side)
    fmt = image_format.upper()
    buf = io.BytesIO()
    if fmt in {"JPG", "JPEG"}:
        pil = pil.convert("RGB")
        pil.save(buf, format="JPEG", quality=quality, optimize=True)
        mime = "image/jpeg"
    elif fmt == "PNG":
        pil.save(buf, format="PNG", optimize=True)
        mime = "image/png"
    else:
        raise ValueError(f"unsupported image_format: {image_format}")
    encoded = base64.b64encode(buf.getvalue()).decode("ascii")
    return f"data:{mime};base64,{encoded}"


def image_sha256(image: Image.Image | bytes) -> str:
    if isinstance(image, bytes):
        payload = image
    else:
        payload = image_to_png_bytes(image)
    return hashlib.sha256(payload).hexdigest()


def resize_for_preview(image: Image.Image, max_side: int = 1400) -> Image.Image:
    pil = load_image(image)
    if max(pil.size) <= max_side:
        return pil
    out = pil.copy()
    out.thumbnail((max_side, max_side), Image.Resampling.LANCZOS)
    return out


def draw_defects(
    image: Image.Image,
    defects: Iterable[dict[str, Any]],
    *,
    title: str | None = None,
    max_boxes: int = 300,
) -> Image.Image:
    """Draw normalized defect boxes onto an RGB copy of an image."""
    out = load_image(image)
    draw = ImageDraw.Draw(out)
    width, height = out.size
    font = ImageFont.load_default()

    if title:
        draw.rectangle((0, 0, min(width, 440), 24), fill=(12, 10, 9))
        draw.text((8, 6), title, fill=(254, 243, 199), font=font)

    drawn = 0
    for defect in defects:
        if drawn >= max_boxes:
            break
        label = str(defect.get("label", "unknown"))
        pixels = bbox_to_pixels(defect.get("bbox"), width, height)
        if pixels is None:
            continue
        x_min, y_min, x_max, y_max = pixels
        color, line_width = LABEL_STYLE.get(label, DEFAULT_STYLE)
        draw.rectangle((x_min, y_min, x_max, y_max), outline=color, width=line_width)

        label_text = LABEL_DISPLAY_NAMES.get(label, label)
        text_bbox = draw.textbbox((x_min, max(0, y_min - 16)), label_text, font=font)
        draw.rectangle(text_bbox, fill=(12, 10, 9))
        draw.text((text_bbox[0] + 1, text_bbox[1]), label_text, fill=color, font=font)
        drawn += 1

    return out


__all__ = [
    "DEFAULT_STYLE",
    "LABEL_STYLE",
    "draw_defects",
    "image_sha256",
    "image_to_data_uri",
    "image_to_png_bytes",
    "load_image",
    "resize_for_preview",
]