backend/app/ml/embedder.py

"""Embedder interface + Mock (default), HF, and fine-tuned re-ID implementations (spec §9.1).

The Embedder is one of the four swap points. The matcher only depends on this Protocol, so a
better model can be dropped in without schema or API changes (A6).
"""
from __future__ import annotations

import hashlib
from typing import Protocol, runtime_checkable

import numpy as np
from PIL import Image

from ..config import settings


@runtime_checkable
class Embedder(Protocol):
    name: str
    version: str
    dim: int

    def embed(self, image_paths: list[str]) -> list[np.ndarray]:
        """Return one L2-normalized float32 vector per input path."""
        ...


def _l2_normalize(v: np.ndarray) -> np.ndarray:
    v = np.asarray(v, dtype=np.float32)
    norm = np.linalg.norm(v)
    if norm == 0:
        # Avoid divide-by-zero; return a stable unit vector.
        out = np.zeros_like(v)
        out[0] = 1.0
        return out
    return v / norm


class MockEmbedder:
    """Deterministic, weight-free embedder (spec §9.1).

    Derives a stable vector from a hash of the *downscaled pixel data*. Same image -> same vector;
    similar-but-not-identical images are NOT meaningfully close. This is intentional: it exists so
    the whole system is runnable and testable without model weights, not to produce real matches.
    """

    name = "mock"
    version = "v1"
    dim = 64

    def __init__(self, dim: int | None = None):
        if dim:
            self.dim = dim

    def _vector_for(self, path: str) -> np.ndarray:
        try:
            with Image.open(path) as img:
                img = img.convert("L").resize((32, 32))
                raw = img.tobytes()
        except Exception:
            # Fall back to hashing the file bytes so the pipeline still produces a vector.
            with open(path, "rb") as fh:
                raw = fh.read()
        # Expand a SHA-256 digest deterministically into `dim` floats.
        seed = int.from_bytes(hashlib.sha256(raw).digest()[:8], "big")
        rng = np.random.default_rng(seed)
        return _l2_normalize(rng.standard_normal(self.dim).astype(np.float32))

    def embed(self, image_paths: list[str]) -> list[np.ndarray]:
        return [self._vector_for(p) for p in image_paths]


class HFEmbedder:  # pragma: no cover - exercised only when transformers is installed
    """HuggingFace image model used as a re-ID embedder via its penultimate pooled features.

    The default model id is a dog-breed classifier whose pre-classifier embeddings are strong for
    individual-dog re-identification (per the project owner). We run the image through the model,
    take the last hidden state, pool it (global average for CNN feature maps, mean-over-tokens for
    transformer sequences), and L2-normalize. Breed *labels* are NOT used here — that is the
    separate BreedClassifier swap point. Lazily imports torch/transformers only when selected.
    """

    def __init__(self, model_id: str | None = None):
        import torch  # noqa: F401
        from transformers import AutoImageProcessor, AutoModelForImageClassification

        self.model_id = model_id or settings.embedder_hf_model
        self.name = "hf-embed"
        # Stable version derived from the repo name so embeddings are tagged per model.
        self.version = self.model_id.rsplit("/", 1)[-1]
        self.processor = AutoImageProcessor.from_pretrained(self.model_id)
        self.model = AutoModelForImageClassification.from_pretrained(self.model_id).eval()
        self._device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model.to(self._device)
        # Infer the embedding dimension with one dummy forward (robust across architectures).
        self.dim = int(self._features(Image.new("RGB", (224, 224))).shape[0])

    @staticmethod
    def _pool(h):
        """Pool the last hidden state to one vector: global-avg for CNN maps, mean for token seqs."""
        if h.dim() == 4:        # [B, C, H, W] CNN feature map -> global average pool
            return h.mean(dim=(2, 3))[0]
        if h.dim() == 3:        # [B, seq, hidden] transformer tokens -> mean pool
            return h.mean(dim=1)[0]
        return h.reshape(h.shape[0], -1)[0]  # already pooled

    def _features(self, img: Image.Image) -> np.ndarray:
        import torch

        with torch.no_grad():
            inputs = self.processor(images=img.convert("RGB"), return_tensors="pt").to(self._device)
            out = self.model(**inputs, output_hidden_states=True)
            return self._pool(out.hidden_states[-1]).float().cpu().numpy()

    def embed(self, image_paths: list[str]) -> list[np.ndarray]:
        out: list[np.ndarray] = []
        for p in image_paths:
            with Image.open(p) as img:
                out.append(_l2_normalize(self._features(img)))
        return out

    def embed_and_breed(
        self, image_paths: list[str], top_k: int
    ) -> list[tuple[np.ndarray, list[tuple[str, float]]]]:
        """One forward pass per image -> ``(embedding_vector, ranked [(label, score)])``.

        This model *is* the breed classifier, so a single forward yields BOTH the penultimate re-ID
        features (embedding) and the softmax breed labels — the model runs exactly ONCE per image.
        The storage pipeline uses this so stored pictures never forward through the model twice.
        """
        import torch

        from .breed import normalize_label

        id2label = self.model.config.id2label
        k = min(top_k, len(id2label))
        results: list[tuple[np.ndarray, list[tuple[str, float]]]] = []
        with torch.no_grad():
            for p in image_paths:
                with Image.open(p) as img:
                    inputs = self.processor(
                        images=img.convert("RGB"), return_tensors="pt"
                    ).to(self._device)
                    out = self.model(**inputs, output_hidden_states=True)
                    vec = _l2_normalize(self._pool(out.hidden_states[-1]).float().cpu().numpy())
                    probs = torch.softmax(out.logits[0], dim=-1)
                    scores, idxs = torch.topk(probs, k)
                    labels = [
                        (normalize_label(id2label[int(i)]), float(s))
                        for s, i in zip(scores.tolist(), idxs.tolist())
                    ]
                    results.append((vec, labels))
        return results


_MEAN, _STD = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]  # ImageNet stats (re-ID preprocessing)


class ReIDEmbedder:  # pragma: no cover - exercised only when torch is installed
    """Fine-tuned re-ID embedder (EMBEDDER=reid): the breed backbone further trained with triplet
    loss for individual-dog re-identification.

    Loads a local checkpoint (a state_dict saved by scripts/train_reid.py's ``ReIDModel`` — keys are
    prefixed ``backbone.``) into the same ResNet base as the breed model, and emits the L2-normalized
    penultimate pooled features. Preprocessing matches TRAINING (Resize 224 + ImageNet norm), NOT the
    HF image processor. This model produces NO breed labels — when it is the active embedder, breed
    labels come from the separate breed classifier via images.py's two-model path.
    """

    def __init__(self, ckpt_path: str | None = None, base_model: str | None = None):
        import torch  # noqa: F401
        import torchvision.transforms as T
        from transformers import AutoModel

        self.name = "reid"
        self.version = settings.reid_model_version
        base = base_model or settings.embedder_hf_model
        self.model = AutoModel.from_pretrained(base)
        raw = torch.load(ckpt_path or settings.reid_model_path, map_location="cpu")
        # accept the ReIDModel wrapper's 'backbone.'-prefixed keys OR a bare backbone state_dict.
        state = {k.removeprefix("backbone."): v for k, v in raw.items()}
        self.model.load_state_dict(state)
        self.model.eval()
        self._device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model.to(self._device)
        self._prep = T.Compose(
            [T.Resize((224, 224)), T.ToTensor(), T.Normalize(_MEAN, _STD)]
        )
        self.dim = 2048

    def embed(self, image_paths: list[str]) -> list[np.ndarray]:
        import torch

        out: list[np.ndarray] = []
        with torch.no_grad():
            for p in image_paths:
                with Image.open(p) as img:
                    x = self._prep(img.convert("RGB")).unsqueeze(0).to(self._device)
                    feat = self.model(x).pooler_output.flatten(1)[0]
                    out.append(_l2_normalize(feat.cpu().numpy()))
        return out


_embedder: Embedder | None = None


def get_embedder() -> Embedder:
    global _embedder
    if _embedder is None:
        if settings.embedder == "hf":
            _embedder = HFEmbedder()
        elif settings.embedder == "reid":
            _embedder = ReIDEmbedder()
        else:
            _embedder = MockEmbedder()
    return _embedder


def reset_embedder_cache() -> None:
    """Test hook to force re-selection after settings change."""
    global _embedder
    _embedder = None