Merge pull request #13 from computational-cell-analytics/box-prompter

Update Prompted Generator - Add Box Prompts

Author: constantinpape

micro_sam/prompt_generators.py

@@ -2,11 +2,17 @@
 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_point_prompts=True, 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
+
+        if self.get_point_prompts is False and self.get_box_prompts is False:
+            raise ValueError("You need to request for box/point prompts or both")
 
     def __call__(self, gt, gt_id, center_coordinates, bbox_coordinates):
         """
@@ -20,9 +26,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
@@ -52,8 +61,10 @@ def __call__(self, gt, gt_id, center_coordinates, bbox_coordinates):
         dilated_object = binary_dilation(object_mask, iterations=self.dilation_strength)
         background_mask = np.zeros(gt.shape)
         background_mask[bbox_coordinates[0]:bbox_coordinates[2], bbox_coordinates[1]:bbox_coordinates[3]] = 1
-        background_mask = abs(background_mask - dilated_object)
         background_mask = binary_dilation(background_mask, iterations=self.dilation_strength)
+        background_mask = abs(
+            background_mask.astype(np.float32) - dilated_object.astype(np.float32)
+        )  # casting booleans to do subtraction
 
         n_negative_remaining = self.n_negative_points
         if n_negative_remaining > 0:
@@ -72,4 +83,12 @@ 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 (TODO: fix it later)
+        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
+
+        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