Adding scores

voxel51 · Oct 5, 2024 · 00c8ffb · 00c8ffb
1 parent 54ce128
commit 00c8ffb
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/fiftyone/utils/sam.py b/fiftyone/utils/sam.py
@@ -271,14 +271,19 @@ def _forward_pass_boxes(self, imgs):
                 device=sam_predictor.device,
             )
 
-            masks, _, _ = sam_predictor.predict_torch(
+            masks, scores, _ = sam_predictor.predict_torch(
                 point_coords=None,
                 point_labels=None,
                 boxes=transformed_boxes,
                 multimask_output=False,
             )
             outputs.append(
-                {"boxes": input_boxes, "labels": labels, "masks": masks}
+                {
+                    "boxes": input_boxes,
+                    "labels": labels,
+                    "masks": masks,
+                    "scores": scores,
+                }
             )
 
         return outputs

diff --git a/fiftyone/utils/sam2.py b/fiftyone/utils/sam2.py
@@ -201,7 +201,7 @@ def _forward_pass_boxes(self, imgs):
                 device=sam2_predictor.device,
             )
 
-            masks, _, _ = sam2_predictor.predict(
+            masks, scores, _ = sam2_predictor.predict(
                 point_coords=None,
                 point_labels=None,
                 box=sam_boxes[None, :],
@@ -214,6 +214,9 @@ def _forward_pass_boxes(self, imgs):
                     "boxes": input_boxes,
                     "labels": labels,
                     "masks": torch.tensor(masks, device=sam2_predictor.device),
+                    "scores": torch.tensor(
+                        scores, device=sam2_predictor.device
+                    ),
                 }
             )