Update Arguments for Multi-Prompting - for Box/Prompt(/both)

Author: anwai98

micro_sam/prompt_generators.py

@@ -3,11 +3,13 @@
 
 
 class PointAndBoxPromptGenerator:
-    def __init__(self, n_positive_points, n_negative_points, dilation_strength, get_box_prompts=False):
+    def __init__(self, n_positive_points, n_negative_points, dilation_strength,
+                 get_point_prompts=False, 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
+        self.get_point_prompts = get_point_prompts
 
     def __call__(self, gt, gt_id, center_coordinates, bbox_coordinates):
         """
@@ -77,7 +79,14 @@ def __call__(self, gt, gt_id, center_coordinates, bbox_coordinates):
                 label_list.append(0)
 
         # returns object-level masks per instance for cross-verification (fix it later)
-        if self.get_box_prompts:
+        if self.get_point_prompts is True and self.get_box_prompts is True:  # we want points and box
             return coord_list, label_list, bbox_list, object_mask
-        else:
+
+        elif self.get_point_prompts is True and self.get_box_prompts is False:  # we want only points
             return coord_list, label_list, None, object_mask
+
+        elif self.get_point_prompts is False and self.get_box_prompts is True:  # we want only boxes
+            return None, None, bbox_list, object_mask
+        else:
+            assert self.get_point_prompts is False and self.get_box_prompts is False, \
+                "You need to request for box/point prompts or both"