adam/eve.py

from __future__ import annotations

from dataclasses import asdict, dataclass
import json
import math
from pathlib import Path
from typing import Callable, Sequence


EVE_MODEL_ID = "facebook/dinov2-small"


@dataclass(slots=True)
class EveResult:
    path: str
    match_score: float
    decision_confidence: float
    suggestion: str


def _normalize(vector: Sequence[float]) -> list[float]:
    length = math.sqrt(sum(float(value) ** 2 for value in vector)) or 1.0
    return [float(value) / length for value in vector]


def _centroid(vectors: Sequence[Sequence[float]]) -> list[float]:
    if not vectors:
        raise ValueError("EVE needs at least one good reference image.")
    width = len(vectors[0])
    if not width or any(len(vector) != width for vector in vectors):
        raise ValueError("EVE received incompatible image embeddings.")
    return _normalize([
        sum(float(vector[index]) for vector in vectors) / len(vectors)
        for index in range(width)
    ])


def _cosine(left: Sequence[float], right: Sequence[float]) -> float:
    return sum(a * b for a, b in zip(_normalize(left), _normalize(right)))


def classify_eve_embeddings(
    image_paths: Sequence[str | Path],
    image_vectors: Sequence[Sequence[float]],
    good_vectors: Sequence[Sequence[float]],
    bad_vectors: Sequence[Sequence[float]] = (),
    *,
    keep_threshold: float = 0.75,
    reject_threshold: float = 0.25,
) -> list[EveResult]:
    """Classify embeddings while keeping the uncertain middle reviewable."""
    if len(image_paths) != len(image_vectors):
        raise ValueError("EVE needs one embedding for every dataset image.")
    if not 0 <= reject_threshold < keep_threshold <= 1:
        raise ValueError("EVE thresholds must leave an uncertain middle range.")
    good_center = _centroid(good_vectors)
    bad_center = _centroid(bad_vectors) if bad_vectors else None

    # With only good references, calibrate the decision boundary from how
    # tightly the references agree with their own centroid.
    reference_floor = min(_cosine(vector, good_center) for vector in good_vectors)
    one_class_center = max(0.20, reference_floor - 0.18)
    results: list[EveResult] = []
    for raw_path, vector in zip(image_paths, image_vectors):
        positive = _cosine(vector, good_center)
        if bad_center is not None:
            negative = _cosine(vector, bad_center)
            # Temperature-scaled two-prototype probability.
            delta = max(-30.0, min(30.0, (positive - negative) / 0.08))
            score = 1.0 / (1.0 + math.exp(-delta))
        else:
            delta = max(-30.0, min(30.0, (positive - one_class_center) / 0.08))
            score = 1.0 / (1.0 + math.exp(-delta))
        if score >= keep_threshold:
            suggestion = "keep"
            confidence = score
        elif score <= reject_threshold:
            suggestion = "reject"
            confidence = 1.0 - score
        else:
            suggestion = "unreviewed"
            confidence = max(score, 1.0 - score)
        results.append(EveResult(str(Path(raw_path).resolve()), score, confidence, suggestion))
    return results


class EveVisionModel:
    """Lazy local DINOv2 feature extractor used by EVE."""

    def __init__(self, model_id: str = EVE_MODEL_ID, *, prefer_gpu: bool = True) -> None:
        self.model_id = model_id
        self.prefer_gpu = prefer_gpu
        self._processor = None
        self._model = None
        self._device = "cpu"

    def load(self) -> None:
        if self._model is not None:
            return
        try:
            import os
            os.environ.setdefault("USE_TF", "0")
            os.environ.setdefault("USE_FLAX", "0")
            import torch
            from transformers import AutoImageProcessor, AutoModel
        except ImportError as exc:
            raise RuntimeError(
                "EVE needs PyTorch and Transformers. Launch ADAM with its normal Python environment."
            ) from exc
        self._device = "cuda" if self.prefer_gpu and torch.cuda.is_available() else "cpu"
        self._processor = AutoImageProcessor.from_pretrained(self.model_id, use_fast=True)
        self._model = AutoModel.from_pretrained(self.model_id).to(self._device).eval()

    def unload(self) -> None:
        self._processor = None
        self._model = None
        try:
            import torch
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        except ImportError:
            pass

    def embed(self, paths: Sequence[str | Path], progress: Callable[[int, int], None] | None = None) -> list[list[float]]:
        self.load()
        import torch
        from PIL import Image

        vectors: list[list[float]] = []
        total = len(paths)
        for index, path in enumerate(paths, 1):
            try:
                with Image.open(path) as source:
                    image = source.convert("RGB")
                    inputs = self._processor(images=image, return_tensors="pt")
                inputs = {key: value.to(self._device) for key, value in inputs.items()}
                with torch.inference_mode():
                    output = self._model(**inputs).last_hidden_state[:, 0]
                    vector = torch.nn.functional.normalize(output, dim=-1)[0].cpu().tolist()
            except Exception as exc:
                raise RuntimeError(f"EVE could not analyze {Path(path).name}: {exc}") from exc
            vectors.append(vector)
            if progress:
                progress(index, total)
        return vectors

    def review(
        self,
        dataset_paths: Sequence[str | Path],
        good_references: Sequence[str | Path],
        bad_references: Sequence[str | Path] = (),
        *,
        keep_threshold: float = 0.75,
        reject_threshold: float = 0.25,
        progress: Callable[[int, int], None] | None = None,
    ) -> list[EveResult]:
        references = [*good_references, *bad_references]
        reference_vectors = self.embed(references)
        image_vectors = self.embed(dataset_paths, progress)
        split = len(good_references)
        return classify_eve_embeddings(
            dataset_paths, image_vectors, reference_vectors[:split], reference_vectors[split:],
            keep_threshold=keep_threshold, reject_threshold=reject_threshold,
        )


def save_eve_results(root: Path, dataset_path: str, results: Sequence[EveResult]) -> Path:
    """Persist the latest EVE proposal so it can be audited after application."""
    import hashlib

    key = hashlib.sha1(str(Path(dataset_path).resolve()).encode("utf-8")).hexdigest()[:12]
    destination = root.resolve() / "data" / "eve_reviews" / f"{key}.json"
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_suffix(".tmp")
    temporary.write_text(json.dumps({
        "agent": "EVE", "dataset_path": str(Path(dataset_path).resolve()),
        "results": [asdict(result) for result in results],
    }, indent=2), encoding="utf-8")
    temporary.replace(destination)
    return destination