src/images.py

"""Image helpers: loading, resizing and cheap (non-AI) augmentation."""

from __future__ import annotations

import hashlib
import random
import re
from pathlib import Path
from typing import List

from PIL import Image, ImageEnhance, ImageOps


def slugify(text: str, max_len: int = 80) -> str:
    text = str(text).strip().lower()
    text = re.sub(r"[^a-z0-9]+", "_", text)
    text = re.sub(r"_+", "_", text).strip("_")
    return (text or "item")[:max_len]


def stable_hash(text: str, length: int = 12) -> str:
    return hashlib.sha1(text.encode("utf-8")).hexdigest()[:length]


def load_rgb(path: str | Path) -> Image.Image:
    return Image.open(path).convert("RGB")


def snap_size(width: int, height: int, longest_side: int) -> tuple[int, int]:
    """Scale so the longest side == longest_side, snapped to multiples of 8."""
    if width >= height:
        new_w = longest_side
        new_h = int(round(longest_side * height / width))
    else:
        new_h = longest_side
        new_w = int(round(longest_side * width / height))
    new_w = max(8, (new_w // 8) * 8)
    new_h = max(8, (new_h // 8) * 8)
    return new_w, new_h


def fit_image(image: Image.Image, longest_side: int) -> Image.Image:
    w, h = image.size
    new_w, new_h = snap_size(w, h, longest_side)
    return image.resize((new_w, new_h), Image.LANCZOS)


def save_jpeg(image: Image.Image, path: Path, quality: int = 92) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    image.convert("RGB").save(path, format="JPEG", quality=quality)


def plain_augment(image: Image.Image, rng: random.Random) -> Image.Image:
    """Light, label-preserving augmentation using only Pillow."""
    out = image
    if rng.random() < 0.5:
        out = ImageOps.mirror(out)
    # Brightness / contrast / colour jitter.
    out = ImageEnhance.Brightness(out).enhance(rng.uniform(0.85, 1.15))
    out = ImageEnhance.Contrast(out).enhance(rng.uniform(0.9, 1.1))
    out = ImageEnhance.Color(out).enhance(rng.uniform(0.9, 1.1))
    # Small rotation with edge-replicating fill.
    angle = rng.uniform(-6.0, 6.0)
    if abs(angle) > 0.5:
        out = out.rotate(angle, resample=Image.BICUBIC, expand=False)
    return out


def read_images_from_dir(directory: str | Path) -> List[Path]:
    exts = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
    return sorted(p for p in Path(directory).iterdir() if p.suffix.lower() in exts)