src/deepfake_scanner/preprocess.py

"""
Image decoding.

In the new architecture, each detector does its own tensor preprocessing
(CLIP needs CLIP's specific resize/normalise; frequency detectors don't
resize at all). This module just turns raw bytes into a validated PIL Image
that all detectors can consume.
"""
from __future__ import annotations

from io import BytesIO

from PIL import Image, UnidentifiedImageError

# Pillow's default decompression-bomb guard is ~89 MP. Tighten it: a 25 MP
# image is more than enough for any consumer photo, and rejecting larger
# inputs caps memory usage on the inference server.
Image.MAX_IMAGE_PIXELS = 25_000_000


class InvalidImageError(ValueError):
    """Raised when the bytes do not decode as a usable image."""


def decode(image_bytes: bytes) -> Image.Image:
    """Decode raw bytes into an RGB PIL image. Raises InvalidImageError on failure."""
    try:
        image = Image.open(BytesIO(image_bytes))
        image.load()  # Force decode now so we catch errors here, not later.
    except (UnidentifiedImageError, OSError) as exc:
        raise InvalidImageError(f"Could not decode image: {exc}") from exc
    except Image.DecompressionBombError as exc:
        raise InvalidImageError(f"Image too large: {exc}") from exc

    if image.mode != "RGB":
        image = image.convert("RGB")

    return image