Update Prompted Generator - Add Box Prompts

Author: anwai98

micro_sam/prompt_generators.py

@@ -2,11 +2,12 @@
 from scipy.ndimage import binary_dilation
 
 
-class PointPromptGenerator:
-    def __init__(self, n_positive_points, n_negative_points, dilation_strength):
+class PointAndBoxPromptGenerator:
+    def __init__(self, n_positive_points, n_negative_points, dilation_strength, get_box_prompts=False):
         self.n_positive_points = n_positive_points
         self.n_negative_points = n_negative_points
         self.dilation_strength = dilation_strength
+        self.get_box_prompts = get_box_prompts
 
     def __call__(self, gt, gt_id, center_coordinates, bbox_coordinates):
         """
@@ -20,9 +21,12 @@ def __call__(self, gt, gt_id, center_coordinates, bbox_coordinates):
         label_list = []
 
         # getting the center coordinate as the first positive point
-        coord_list.append(tuple(map(int, center_coordinates)))
+        coord_list.append(tuple(map(int, center_coordinates)))  # to get int coords instead of float
         label_list.append(1)
 
+        if self.get_box_prompts:
+            bbox_list = [bbox_coordinates]
+
         object_mask = gt == gt_id + 1  # alloting a label id to obtain the coordinates of desired seeds
 
         # getting the additional positive points by randomly sampling points from this mask
@@ -72,4 +76,8 @@ def __call__(self, gt, gt_id, center_coordinates, bbox_coordinates):
                 coord_list.append(negative_coordinates)
                 label_list.append(0)
 
-        return coord_list, label_list, None, object_mask
+        # returns object-level masks per instance for cross-verification (fix it later)
+        if self.get_box_prompts:
+            return coord_list, label_list, bbox_list, object_mask
+        else:
+            return coord_list, label_list, None, object_mask