feat: owlv2 extra bounding box filtering

tests/models/test_owlv2.py

@@ -20,6 +20,27 @@ def test_successful_image_detection():
         assert pred.label == "a photo of a cat"
 
 
+def test_successful_removing_extra_bbox():
+    test_image = "eggs-food-easter-food-drink-44c10e-1024.jpg"
+    prompts = ["egg"]
+
+    image = Image.open(f"tests/shared_data/images/{test_image}")
+
+    owlv2 = Owlv2()
+
+    results = owlv2(prompts=prompts, image=image)
+
+    assert len(results[0]) > 0
+
+    bboxlabels = results[0]
+
+    for bbox_label in bboxlabels:
+        assert bbox_label.label == "egg"
+    current_count = len(bboxlabels)
+    expected_max_count = 42
+    assert current_count <= expected_max_count
+
+
 def test_successful_video_detection():
     test_video = "test_video_5_frames.mp4"
     file_path = f"tests/shared_data/videos/{test_video}"

vision_agent_tools/models/owlv2.py

@@ -1,14 +1,16 @@
-from typing import Optional
+from typing import List, Optional, Tuple, Union
+
 import numpy as np
 import torch
 from PIL import Image
-from typing import List, Tuple, Union
 from pydantic import BaseModel, Field
 from transformers import Owlv2ForObjectDetection, Owlv2Processor
-from transformers.utils import TensorType
 from transformers.image_transforms import center_to_corners_format
 from transformers.models.owlv2.image_processing_owlv2 import box_iou
-from vision_agent_tools.shared_types import BaseMLModel, Device, VideoNumpy, BboxLabel
+from transformers.utils import TensorType
+
+from vision_agent_tools.models.utils import filter_redundant_boxes
+from vision_agent_tools.shared_types import BaseMLModel, BboxLabel, Device, VideoNumpy
 
 
 class OWLV2Config(BaseModel):
@@ -54,6 +56,37 @@ class Owlv2(BaseMLModel):
     and bounding boxes for detected objects with confidence exceeding a threshold.
     """
 
+    from typing import Dict, List
+
+    def _filter_bboxes(self, bboxlabels: List[BboxLabel]) -> List[BboxLabel]:
+        """
+        Filters out redundant BboxLabel objects that fully contain multiple smaller boxes of the same label.
+
+        Parameters:
+            bboxlabels (List[BboxLabel]): List of BboxLabel objects to be filtered.
+
+        Returns:
+            List[BboxLabel]: Filtered list of BboxLabel objects.
+        """
+        bboxes = [bl.bbox for bl in bboxlabels]
+        labels = [bl.label for bl in bboxlabels]
+
+        filtered = filter_redundant_boxes({"bboxes": bboxes, "labels": labels})
+        filtered_bboxes = filtered["bboxes"]
+        filtered_labels = filtered["labels"]
+
+        filtered_pairs = list(zip(filtered_bboxes, filtered_labels))
+
+        # preserving the original order
+        output_bboxlabels = []
+        for bl in bboxlabels:
+            pair = (bl.bbox, bl.label)
+            if pair in filtered_pairs:
+                output_bboxlabels.append(bl)
+                filtered_pairs.remove(pair)  # Remove to handle duplicates correctly
+
+        return output_bboxlabels
+
     def __run_inference(
         self, image, texts, confidence, nms_threshold
     ) -> list[BboxLabel]:
@@ -91,7 +124,9 @@ def __run_inference(
                 BboxLabel(label=texts[i][label.item()], score=score.item(), bbox=box)
             )
 
-        return inferences
+        filtered_inferences = self._filter_bboxes(inferences)
+
+        return filtered_inferences
 
     def __init__(self, model_config: Optional[OWLV2Config] = None):
         """
@@ -247,10 +282,18 @@ def post_process_object_detection_with_nms(
             boxes = boxes * scale_fct[:, None, :]
 
         results = []
-        for s, l, b in zip(scores, labels, boxes):
-            score = s[s > threshold]
-            label = l[s > threshold]
-            box = b[s > threshold]
-            results.append({"scores": score, "labels": label, "boxes": box})
+        for score_array, label_array, box_array in zip(scores, labels, boxes):
+            high_score_mask = score_array > threshold
+            filtered_scores = score_array[high_score_mask]
+            filtered_labels = label_array[high_score_mask]
+            filtered_boxes = box_array[high_score_mask]
+
+            results.append(
+                {
+                    "scores": filtered_scores,
+                    "labels": filtered_labels,
+                    "boxes": filtered_boxes,
+                }
+            )
 
         return results