Vase-Matcher/pipeline/vase_compare.py

"""
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,
    }