feat(ml): argmax grounding in score_image → suggestions carry the winning crop (#133 step 1)

score_image now keeps the ARGMAX beside the max-over-bag: which bag row won each
head. The region query also selects bbox/kind/detector_version, a parallel
bag_meta maps each row → its region (None for the whole-image vector), and every
hit gains grounding {bbox,kind,detector} (null when the global vector won). Threaded
through SuggestionService (new Suggestion.grounding field) → /api/.../suggestions
payload. This is the data the #1206 hover-overlay draws. CCIP-only hits ground null
for now (figure grounding = step 2). Tests: winning crop grounds the tag with its
bbox+kind; whole-image win → grounding None.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01CDgx8bQS5YrGRK76v8HUnM

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
                 ]

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

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:

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