Add overriding of GroundingDINO bboxes with calculated masks bboxes

Youho99 · Apr 17, 2024 · 5b1c390 · 5b1c390
1 parent 12eea0d
commit 5b1c390
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/autodistill_grounded_sam/grounded_sam.py b/autodistill_grounded_sam/grounded_sam.py
@@ -21,6 +21,8 @@
 
 from autodistill.detection import CaptionOntology, DetectionBaseModel
 
+from supervision.detection.utils import mask_to_xyxy
+
 HOME = os.path.expanduser("~")
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -77,5 +79,8 @@ def predict(self, input: Any) -> sv.Detections:
 
         detections.mask = np.array(result_masks)
 
+        # override GroundingDINO bboxes with calculated masks bboxes
+        detections.xyxy = mask_to_xyxy(detections.mask)
+
         # separate in supervision to combine detections and override class_ids
         return detections