diff --git a/backend/app/api/suggestions.py b/backend/app/api/suggestions.py index cf27125..28df98e 100644 --- a/backend/app/api/suggestions.py +++ b/backend/app/api/suggestions.py @@ -46,6 +46,10 @@ async def get_suggestions(image_id: int): # (not dropped) so the rail can show it rejected + offer # one-click un-reject. "rejected": s.rejected, + # the crop region that produced this tag (#1206) — + # {bbox,kind,detector} or null (whole-image won). Drives + # the hover→overlay highlight. + "grounding": s.grounding, } for s in items ] diff --git a/backend/app/services/ml/heads.py b/backend/app/services/ml/heads.py index 1f4d01a..aad0226 100644 --- a/backend/app/services/ml/heads.py +++ b/backend/app/services/ml/heads.py @@ -375,21 +375,33 @@ async def score_image( # skipped rather than scored by heads trained in a different space; a legacy # NULL version is treated as current (those predate per-row stamping). bag = [] + # Parallel to `bag`: what each row IS, so a surfaced tag can point back at the + # crop that produced it (#1206 grounding). None = the whole-image vector (not + # localized); a dict = a region's {bbox, kind, detector}. + bag_meta: list[dict | None] = [] if img.siglip_embedding is not None and img.siglip_model_version in ( cur_version, None, ): bag.append(np.asarray(img.siglip_embedding, dtype=np.float32)) - region_vecs = ( + bag_meta.append(None) + region_rows = ( await session.execute( - select(ImageRegion.siglip_embedding) + select( + ImageRegion.siglip_embedding, + ImageRegion.rx, ImageRegion.ry, ImageRegion.rw, ImageRegion.rh, + ImageRegion.kind, ImageRegion.detector_version, + ) .where(ImageRegion.image_record_id == image_id) .where(ImageRegion.siglip_embedding.is_not(None)) .where(ImageRegion.embedding_version == cur_version) ) ).all() - for (vec,) in region_vecs: + for vec, rx, ry, rw, rh, kind, detector in region_rows: if vec is not None: bag.append(np.asarray(vec, dtype=np.float32)) + bag_meta.append( + {"bbox": [rx, ry, rw, rh], "kind": kind, "detector": detector} + ) if not bag: return [] @@ -398,7 +410,11 @@ async def score_image( norms[norms == 0] = 1.0 Xn = X / norms Z = Xn @ heads["W"].T + heads["b"] # (B, H) - probs = (1.0 / (1.0 + np.exp(-Z))).max(axis=0) # (H,) best over the bag + probs_bag = 1.0 / (1.0 + np.exp(-Z)) # (B, H) + probs = probs_bag.max(axis=0) # (H,) best over the bag + # ARGMAX beside the max: WHICH bag row won each head → the region that grounds + # the tag (bag_meta[win]); None when the whole-image vector won (#1206). + winners = probs_bag.argmax(axis=0) # (H,) out = [] for i, p in enumerate(probs): if threshold_override is not None: @@ -416,6 +432,7 @@ async def score_image( "name": m["name"], "category": m["category"], "score": float(p), + "grounding": bag_meta[int(winners[i])], }) out.sort(key=lambda d: d["score"], reverse=True) return out diff --git a/backend/app/services/ml/suggestions.py b/backend/app/services/ml/suggestions.py index 71fbd24..772074c 100644 --- a/backend/app/services/ml/suggestions.py +++ b/backend/app/services/ml/suggestions.py @@ -43,6 +43,11 @@ class Suggestion: # the rejection is VISIBLE and REVERSIBLE in the rail (misclick recovery, # operator-asked 2026-06-27) instead of silently vanishing or re-suggesting. rejected: bool = False + # grounding = the crop region that produced this suggestion (#1206): + # {bbox:[x,y,w,h] normalized, kind, detector}. None when the whole-image + # vector won (not localized) or for a CCIP-only hit (figure grounding TBD). + # Lets the rail highlight the exact region on hover. + grounding: dict | None = None @dataclass @@ -103,6 +108,7 @@ class SuggestionService: for h in hits: merged[(h["category"], h["tag_id"])] = { "name": h["name"], "score": h["score"], "source": "head", + "grounding": h.get("grounding"), } for c in ccip_hits: key = ("character", c["tag_id"]) @@ -128,6 +134,7 @@ class SuggestionService: source=m["source"], creates_new_tag=False, rejected=tag_id in rejected, + grounding=m.get("grounding"), ) ) for cat in result.by_category: diff --git a/tests/test_ml_suggestions.py b/tests/test_ml_suggestions.py index f8f4d00..cf2b732 100644 --- a/tests/test_ml_suggestions.py +++ b/tests/test_ml_suggestions.py @@ -187,7 +187,31 @@ async def test_concept_region_surfaces_via_max_over_bag(db): )) await db.commit() general = (await SuggestionService(db).for_image(img.id)).by_category["general"] - assert any(s.canonical_tag_id == tag.id and s.score > 0.7 for s in general) + s = next(x for x in general if x.canonical_tag_id == tag.id) + assert s.score > 0.7 + # #1206: the winning crop grounds the tag — hover highlights THIS region + # (the matching-version crop at 0.4,0.4,0.3,0.3), not the whole image. + assert s.grounding is not None + assert s.grounding["bbox"] == pytest.approx([0.4, 0.4, 0.3, 0.3]) + assert s.grounding["kind"] == "concept" + + +@pytest.mark.asyncio +async def test_grounding_none_when_whole_image_wins(db): + # #1206: when the whole-image vector (not a crop) produces the winning score, + # grounding is None — the tag came from the global vector, not a region. + tag = await TagService(db).find_or_create("sky", TagKind.general) + img = await _img(db, "d1" * 32, _emb(0)) # whole-image ALIGNED w/ head + await _head(db, tag.id, slot=0, suggest_threshold=0.5) + db.add(ImageRegion( # an orthogonal crop (0.5) + image_record_id=img.id, kind="concept", + rx=0.1, ry=0.1, rw=0.2, rh=0.2, + siglip_embedding=_emb(5), embedding_version=await _embver(db), + )) + await db.commit() + general = (await SuggestionService(db).for_image(img.id)).by_category["general"] + s = next(x for x in general if x.canonical_tag_id == tag.id) + assert s.grounding is None @pytest.mark.asyncio