File size: 6,736 Bytes
e0265b9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
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) -> None:
        self.model_id = model_id
        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 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 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