image_sources.py

"""
Load raster images into PIL RGB for captioning and PNG export.

Supports common formats via Pillow; Canon CR2 via rawpy; OpenEXR / Radiance HDR
via OpenCV (preferred) or imageio with simple tonemapping for float HDR.
"""

from __future__ import annotations

from pathlib import Path
from typing import List, Union

import numpy as np
from PIL import Image

INPUT_GLOB_PATTERNS: List[str] = [
    "*.jpg",
    "*.jpeg",
    "*.png",
    "*.JPG",
    "*.JPEG",
    "*.PNG",
    "*.cr2",
    "*.CR2",
    "*.exr",
    "*.EXR",
    "*.hdr",
    "*.HDR",
]


def _tonemap_float_rgb(arr: np.ndarray) -> Image.Image:
    """Map linear / HDR float arrays (H, W, 3) to 8-bit sRGB PIL Image."""
    arr = np.asarray(arr, dtype=np.float64)
    if arr.ndim == 2:
        arr = np.stack([arr, arr, arr], axis=-1)
    if arr.shape[-1] >= 4:
        arr = arr[..., :3]
    arr = np.nan_to_num(arr, nan=0.0, posinf=1e10, neginf=0.0)
    flat = arr.reshape(-1)
    lo, hi = np.percentile(flat, (0.5, 99.5))
    hi = max(hi, lo + 1e-6)
    arr = (arr - lo) / (hi - lo)
    arr = np.clip(arr, 0, 1)
    arr = (arr ** (1.0 / 2.2) * 255.0).astype(np.uint8)
    return Image.fromarray(arr, mode="RGB")


def _load_exr_hdr(path: Path) -> Image.Image:
    suffix = path.suffix.lower()
    errors: list[str] = []

    try:
        import cv2

        flags = cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH
        bgr = cv2.imread(str(path), flags)
        if bgr is not None and bgr.size > 0:
            if bgr.ndim == 2:
                bgr = cv2.cvtColor(bgr, cv2.COLOR_GRAY2BGR)
            if bgr.dtype in (np.float32, np.float64):
                rgb = cv2.cvtColor(bgr.astype(np.float32), cv2.COLOR_BGR2RGB)
                return _tonemap_float_rgb(rgb)
            if bgr.dtype == np.uint16:
                rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
                return _tonemap_float_rgb(rgb.astype(np.float64) / 65535.0)
            rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
            return Image.fromarray(rgb, mode="RGB")
    except ImportError:
        errors.append("opencv-python-headless not installed")
    except Exception as e:
        errors.append(f"OpenCV: {e}")

    try:
        import imageio.v3 as iio

        arr = np.asarray(iio.imread(path))
        if arr.ndim == 2:
            arr = np.stack([arr, arr, arr], axis=-1)
        if arr.shape[-1] >= 4:
            arr = arr[..., :3]
        if arr.dtype in (np.float32, np.float64):
            return _tonemap_float_rgb(arr)
        if arr.dtype == np.uint16:
            return _tonemap_float_rgb(arr.astype(np.float64) / 65535.0)
        if arr.dtype != np.uint8 and arr.size and float(arr.max()) <= 1.0:
            arr = (np.clip(arr, 0, 1) * 255.0).astype(np.uint8)
        return Image.fromarray(arr, mode="RGB")
    except ImportError:
        errors.append("imageio not installed")
    except Exception as e:
        errors.append(f"imageio: {e}")

    try:
        img = Image.open(path)
        return img.convert("RGB")
    except Exception as e:
        errors.append(f"Pillow: {e}")

    hint = (
        f"Could not read {path.name} ({suffix}). "
        "Install opencv-python-headless and/or imageio (pip). Details: "
        + "; ".join(errors)
    )
    raise ValueError(hint)


def load_image_rgb(path: Union[str, Path]) -> Image.Image:
    """Open any supported file and return an RGB PIL Image."""
    path = Path(path)
    suffix = path.suffix.lower()

    if suffix == ".cr2":
        try:
            import rawpy
        except ImportError as e:
            raise ImportError(
                "Canon CR2 support requires rawpy. Install: pip install rawpy"
            ) from e
        with rawpy.imread(str(path)) as raw:
            rgb = raw.postprocess(
                use_camera_wb=True,
                no_auto_bright=False,
                output_bps=8,
            )
        return Image.fromarray(rgb, "RGB")

    if suffix in (".exr", ".hdr"):
        return _load_exr_hdr(path)

    img = Image.open(path)
    return img.convert("RGB")