File size: 9,180 Bytes
b5442b0 de1e3fc b5442b0 de1e3fc b5442b0 de1e3fc b5442b0 de1e3fc b5442b0 de1e3fc | 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 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 | """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
|