pipeline/preprocessing.py

"""
Preprocessing — decode, resize, color-convert.

Uses cores.vision for all image operations — no duplicated decode,
resize, or format-sniffing logic.  Preserves the original bytes so
metadata / forensics providers can use them.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import numpy as np

from cores.vision import bytes_to_numpy, resize_with_aspect, sniff_format


@dataclass
class PreprocessedImage:
    """Output of preprocessing — the canonical image object passed downstream."""
    image: np.ndarray        # BGR uint8 HxWx3
    width: int
    height: int
    channels: int = 3
    source: str = ""         # "url" | "base64" | "bytes"
    resized: bool = False
    original_bytes: Optional[bytes] = None
    original_format: Optional[str] = None


class ImagePreprocessor:
    """Decodes + resizes inbound images."""

    def __init__(self, max_dim: int = 1024) -> None:
        self._max_dim = max_dim

    def from_bytes(self, data: bytes, source: str = "bytes") -> PreprocessedImage:
        img = bytes_to_numpy(data)
        fmt = sniff_format(data)
        return self._finalize(img, source, original_bytes=data, original_format=fmt)

    def from_url(self, url: str, timeout: int = 15) -> PreprocessedImage:
        from cores.vision import url_to_bytes
        data = url_to_bytes(url, timeout=timeout)
        img = bytes_to_numpy(data)
        fmt = sniff_format(data)
        return self._finalize(img, "url", original_bytes=data, original_format=fmt)

    def from_numpy(self, img: np.ndarray, source: str = "in_memory") -> PreprocessedImage:
        return self._finalize(img, source)

    def _finalize(
        self,
        img: np.ndarray,
        source: str,
        original_bytes: Optional[bytes] = None,
        original_format: Optional[str] = None,
    ) -> PreprocessedImage:
        # Ensure 3 channels
        if img.ndim == 2:
            import cv2
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        elif img.shape[2] == 4:
            import cv2
            img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)

        resized = False
        h, w = img.shape[:2]
        if max(h, w) > self._max_dim:
            img = resize_with_aspect(img, max_dim=self._max_dim)
            resized = True

        h, w = img.shape[:2]
        return PreprocessedImage(
            image=img,
            width=w,
            height=h,
            channels=img.shape[2] if img.ndim == 3 else 1,
            source=source,
            resized=resized,
            original_bytes=original_bytes,
            original_format=original_format,
        )