fix(server): map GroundingDINO detections back onto the caller's labels (#277)

* fix(server): map GroundingDINO detections back onto the caller's labels

The post-processor returns lowercased phrase fragments, so a caller matching
detections against the labels it sent found none. Each decoded phrase is now
mapped onto the caller's label it overlaps most; a free-text instruction
prompt keeps the phrase.

* fix(server): count each distinct label word once when canonicalising

Author: krisztian-gajdar

packages/sie_server/src/sie_server/adapters/grounding_dino/adapter.py

@@ -28,6 +28,7 @@
 from __future__ import annotations
 
 import logging
+import re
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar
 
@@ -53,6 +54,34 @@
 _ERR_ENCODE_NOT_SUPPORTED = "GroundingDINOAdapter does not support encode(). Use extract() instead."
 
 
+_WORD_RE = re.compile(r"\w+")
+
+
+def _canonical_label(phrase: str, labels: list[str] | None) -> str:
+    """The caller's label a grounded phrase came from, or the phrase itself.
+
+    Scores each label by the share of its words the phrase contains, then by
+    how many of the phrase's words it explains ("red handbag" beats "handbag"
+    for the phrase "red handbag"); the earlier label breaks a full tie. The
+    phrase is kept verbatim only when no label shares a word with it, which
+    is the ``instruction`` prompt case (no labels) or a decode the caller
+    never asked for.
+    """
+    if not labels:
+        return phrase
+    phrase_words = set(_WORD_RE.findall(phrase.lower()))
+    best_label, best_key = phrase, (0.0, 0)
+    for label in labels:
+        words = set(_WORD_RE.findall(label.lower()))
+        if not words:
+            continue
+        matched = len(words & phrase_words)
+        key = (matched / len(words), matched)
+        if key > best_key:
+            best_label, best_key = label, key
+    return best_label
+
+
 class GroundingDINOAdapter(BaseAdapter):
     """Adapter for GroundingDINO open-vocabulary object detection.
 
@@ -225,6 +254,7 @@ def extract(
                 pixel_values=pixel_values,
                 pixel_mask=pixel_mask,
                 original_sizes=original_sizes,
+                labels=labels,
             )
         else:
             # Fallback: decode images inline
@@ -246,6 +276,7 @@ def extract(
                 box_threshold=box_threshold,
                 text_threshold=text_threshold,
                 images=images,
+                labels=labels,
             )
 
         # Map results back to original item positions
@@ -265,6 +296,7 @@ def _detect_batch(
         pixel_mask: torch.Tensor | None = None,
         original_sizes: list[tuple[int, int]] | None = None,
         images: list[Image] | None = None,
+        labels: list[str] | None = None,
     ) -> list[list[DetectedObject]]:
         """Run batched detection on multiple images.
 
@@ -279,6 +311,8 @@ def _detect_batch(
             pixel_values: Preprocessed image tensor [B, C, H, W] (optional).
             original_sizes: List of (width, height) tuples for bbox denormalization.
             images: List of PIL Images (fallback if pixel_values not provided).
+            labels: The caller's labels, verbatim, when the prompt was built from
+                them; every detection's label is mapped back onto one of these.
 
         Returns:
             List of detected object lists, one per input image.
@@ -345,10 +379,19 @@ def _detect_batch(
         )
 
         # Convert results to DetectedObject format
-        return [self._results_to_objects(result) for result in results]
-
-    def _results_to_objects(self, result: dict[str, Any]) -> list[DetectedObject]:
-        """Convert post-processed detection result to DetectedObject list."""
+        return [self._results_to_objects(result, labels) for result in results]
+
+    def _results_to_objects(self, result: dict[str, Any], labels: list[str] | None = None) -> list[DetectedObject]:
+        """Convert post-processed detection result to DetectedObject list.
+
+        The post-processor decodes, per box, the prompt tokens above
+        ``text_threshold``, so a caller's label comes back lowercased and
+        possibly as a fragment or a span merged across two labels ("leather
+        handbag" for "Red Leather Handbag"). Callers match detections against
+        the labels they sent, so each phrase is mapped back onto the caller's
+        label it overlaps most; a free-text ``instruction`` prompt has no
+        labels and keeps the phrase.
+        """
         boxes = result["boxes"]
         scores = result["scores"]
         result_labels = result.get("text_labels", result.get("labels", []))
@@ -364,7 +407,7 @@ def _results_to_objects(self, result: dict[str, Any]) -> list[DetectedObject]:
         for i in range(n_detections):
             x1, y1, x2, y2 = boxes_list[i]
             score = scores_list[i]
-            label_text = result_labels[i]
+            label_text = _canonical_label(result_labels[i], labels)
 
             objects.append(
                 DetectedObject(

packages/sie_server/tests/adapters/test_grounding_dino.py

@@ -9,7 +9,7 @@
 import pytest
 import torch
 from PIL import Image
-from sie_server.adapters.grounding_dino.adapter import GroundingDINOAdapter
+from sie_server.adapters.grounding_dino.adapter import GroundingDINOAdapter, _canonical_label
 from sie_server.core.inference_output import ExtractOutput
 from sie_server.types.inputs import ImageInput, Item
 
@@ -319,6 +319,44 @@ def test_results_to_objects_bulk_converts_tensors_once(self) -> None:
         scores.detach.return_value.cpu.assert_called_once_with()
         scores.detach.return_value.cpu.return_value.tolist.assert_called_once_with()
 
+    def test_results_to_objects_maps_decoded_phrases_back_onto_the_callers_labels(self) -> None:
+        """The post-processor lowercases and fragments multi-word labels; callers get their own back."""
+        adapter = GroundingDINOAdapter("IDEA-Research/grounding-dino-tiny")
+        boxes = MagicMock()
+        boxes.__len__.return_value = 3
+        boxes.detach.return_value.cpu.return_value.tolist.return_value = [
+            [10.0, 20.0, 110.0, 220.0],
+            [0.0, 0.0, 50.0, 50.0],
+            [5.0, 5.0, 15.0, 15.0],
+        ]
+        scores = MagicMock()
+        scores.detach.return_value.cpu.return_value.tolist.return_value = [0.9, 0.5, 0.4]
+
+        objects = adapter._results_to_objects(
+            {
+                "boxes": boxes,
+                "scores": scores,
+                "text_labels": ["leather handbag", "handbag backpack", "camera"],
+            },
+            ["Red Leather Handbag", "backpack", "camera — acceptance 1959"],
+        )
+
+        # A span merged across two labels ("handbag backpack") goes to the
+        # label it covers completely, not to the one it merely touches.
+        assert [obj["label"] for obj in objects] == ["Red Leather Handbag", "backpack", "camera — acceptance 1959"]
+
+    def test_canonical_label_prefers_the_label_the_phrase_covers_most(self) -> None:
+        labels = ["handbag", "red handbag", "backpack"]
+        assert _canonical_label("handbag", labels) == "handbag"
+        assert _canonical_label("red handbag", labels) == "red handbag"
+        assert _canonical_label("backpack.", labels) == "backpack"
+        # A repeated word counts once, so it cannot outscore the earlier label.
+        assert _canonical_label("red car", ["red car", "red red car"]) == "red car"
+        # No labels (an instruction prompt) or no overlap: the phrase stands.
+        assert _canonical_label("dog", None) == "dog"
+        assert _canonical_label("dog", labels) == "dog"
+        assert _canonical_label("dog", ["", "  "]) == "dog"
+
     def test_results_to_objects_empty_does_not_transfer_scores(self) -> None:
         adapter = GroundingDINOAdapter("IDEA-Research/grounding-dino-tiny")
         scores = MagicMock()