Add project files

pipeline/vase_compare.py

@@ -0,0 +1,158 @@
+"""
+Vase-identity comparison: crops the vase region (via SAM3's "vase" concept
+mask -- see pipeline/sam3_client.py) from both the upload and the matched
+template, then compares the two crops with two complementary embedding
+models:
+
+  - DINOv2 (facebook/dinov2-base): self-supervised, patch-level visual
+    features -- good at fine-grained shape/texture/material detail.
+  - CLIP (openai/clip-vit-base-patch32): contrastive image embedding -- a
+    coarser, more holistic notion of visual similarity, used as a second
+    opinion that isn't fooled by the same quirks DINO might be.
+
+Purely informational, opt-in (run alongside the SAM/YOLO-World flower count,
+triggered by the same button) -- never feeds into matching/scoring.
+"""
+
+import gc
+import logging
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+
+import config
+
+logger = logging.getLogger(__name__)
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+_dino_model = None
+_dino_processor = None
+_clip_model = None
+_clip_processor = None
+
+
+def _get_dino():
+    global _dino_model, _dino_processor
+    if _dino_model is None:
+        logger.info("Loading DINOv2 (%s) on %s...", config.DINO_MODEL_NAME, DEVICE)
+        from transformers import AutoImageProcessor, AutoModel
+        _dino_processor = AutoImageProcessor.from_pretrained(config.DINO_MODEL_NAME)
+        _dino_model = AutoModel.from_pretrained(config.DINO_MODEL_NAME).eval().to(DEVICE)
+    return _dino_model, _dino_processor
+
+
+def _get_clip():
+    global _clip_model, _clip_processor
+    if _clip_model is None:
+        logger.info("Loading CLIP (%s) on %s...", config.CLIP_MODEL_NAME, DEVICE)
+        from transformers import CLIPModel, CLIPProcessor
+        _clip_model = CLIPModel.from_pretrained(config.CLIP_MODEL_NAME).eval().to(DEVICE)
+        _clip_processor = CLIPProcessor.from_pretrained(config.CLIP_MODEL_NAME)
+    return _clip_model, _clip_processor
+
+
+def unload_models():
+    global _dino_model, _dino_processor, _clip_model, _clip_processor
+    freed = _dino_model is not None or _clip_model is not None
+    _dino_model = None
+    _dino_processor = None
+    _clip_model = None
+    _clip_processor = None
+    if freed:
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    return freed
+
+
+def _to_pil(bgr):
+    return Image.fromarray(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
+
+
+def _dino_embed(bgr):
+    model, processor = _get_dino()
+    with torch.no_grad():
+        inputs = processor(images=_to_pil(bgr), return_tensors="pt").to(DEVICE)
+        out = model(**inputs)
+        feat = out.last_hidden_state[:, 0]  # CLS token
+        feat = torch.nn.functional.normalize(feat, dim=-1)
+    return feat.cpu().numpy()[0]
+
+
+def _clip_embed(bgr):
+    model, processor = _get_clip()
+    with torch.no_grad():
+        inputs = processor(images=_to_pil(bgr), return_tensors="pt").to(DEVICE)
+        feat = model.get_image_features(**inputs)
+        feat = torch.nn.functional.normalize(feat, dim=-1)
+    return feat.cpu().numpy()[0]
+
+
+def _cosine_pct(a, b):
+    """Cosine similarity of two L2-normalized vectors, clamped to [0, 1] and
+    reported as a percentage -- negative similarity (near-opposite vectors)
+    is clamped to 0% rather than reported as a negative number, since
+    "how similar" isn't meaningful past that point for this use case."""
+    sim = float(np.dot(a, b))
+    return round(max(0.0, min(1.0, sim)) * 100, 1)
+
+
+def crop_mask(bgr, mask_bool):
+    """Crops the bounding box of mask_bool out of bgr, with a small pad
+    (config.VASE_CROP_PAD_FRAC) so the mask's own (occasionally imprecise)
+    edge doesn't cut off the vase's rim or base -- then blacks out any
+    pixel inside that padded box that the mask doesn't cover, so DINOv2/CLIP
+    only ever see vase pixels, not whatever flower stems or ribbon happen to
+    share the box's corners. Returns None if the mask is empty."""
+    ys, xs = np.nonzero(mask_bool)
+    if len(xs) == 0:
+        return None
+    h, w = bgr.shape[:2]
+    x1, y1, x2, y2 = xs.min(), ys.min(), xs.max(), ys.max()
+    bw, bh = x2 - x1, y2 - y1
+    pad = config.VASE_CROP_PAD_FRAC
+    px1, px2 = x1 - bw * pad, x2 + bw * pad
+    py1, py2 = y1 - bh * pad, y2 + bh * pad
+    px1, py1 = max(0, int(round(px1))), max(0, int(round(py1)))
+    px2, py2 = min(w, int(round(px2))), min(h, int(round(py2)))
+    if px2 <= px1 or py2 <= py1:
+        return None
+
+    crop = bgr[py1:py2, px1:px2].copy()
+    mask_crop = mask_bool[py1:py2, px1:px2]
+    crop[~mask_crop] = 0
+    return crop
+
+
+def clip_similarity_pct(crop_a, crop_b):
+    """CLIP-only similarity between two crops -- no DINOv2. Used for the
+    auto-triggered "flower check" beside the weighted verdict (see
+    engine.flower_summary): unlike the opt-in vase comparison, that one runs
+    on every confident match, so it deliberately loads only CLIP (small,
+    fast) rather than both embedding models."""
+    return _cosine_pct(_clip_embed(crop_a), _clip_embed(crop_b))
+
+
+def compare_vases(input_crop_bgr, template_crop_bgr):
+    dino_pct = _cosine_pct(_dino_embed(input_crop_bgr), _dino_embed(template_crop_bgr))
+    clip_pct = _cosine_pct(_clip_embed(input_crop_bgr), _clip_embed(template_crop_bgr))
+    combined_pct = round(
+        config.VASE_DINO_WEIGHT * dino_pct + config.VASE_CLIP_WEIGHT * clip_pct, 1
+    )
+
+    if combined_pct >= config.VASE_SAME_THRESHOLD:
+        verdict = "same"
+    elif combined_pct >= config.VASE_UNCERTAIN_THRESHOLD:
+        verdict = "uncertain"
+    else:
+        verdict = "different"
+
+    return {
+        "dino_similarity_pct": dino_pct,
+        "clip_similarity_pct": clip_pct,
+        "combined_pct": combined_pct,
+        "verdict": verdict,
+    }