Merge pull request #61 from computational-cell-analytics/update-prompter

Update Prompt Generator - First Point as Any FG Point

Author: constantinpape

micro_sam/prompt_generators.py

@@ -14,7 +14,7 @@ def __init__(self, n_positive_points, n_negative_points, dilation_strength,
         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):
+    def __call__(self, gt, gt_id, bbox_coordinates, center_coordinates=None):
         """
         Parameters:
             gt: True Labels
@@ -25,17 +25,23 @@ def __call__(self, gt, gt_id, center_coordinates, bbox_coordinates):
         coord_list = []
         label_list = []
 
-        # getting the center coordinate as the first positive point
-        coord_list.append(tuple(map(int, center_coordinates)))  # to get int coords instead of float
-        label_list.append(1)
+        if center_coordinates is not None:
+            # getting the center coordinate as the first positive point (OPTIONAL)
+            coord_list.append(tuple(map(int, center_coordinates)))  # to get int coords instead of float
+            label_list.append(1)
+
+            # getting the additional positive points by randomly sampling points from this mask except the center coordinate
+            n_positive_remaining = self.n_positive_points - 1
+
+        else:
+            # need to sample "self.n_positive_points" number of points
+            n_positive_remaining = self.n_positive_points
 
         if self.get_box_prompts:
             bbox_list = [bbox_coordinates]
 
         object_mask = gt == gt_id
 
-        # getting the additional positive points by randomly sampling points from this mask
-        n_positive_remaining = self.n_positive_points - 1
         if n_positive_remaining > 0:
             # all coordinates of our current object
             object_coordinates = np.where(object_mask)

test/test_prompt_generators.py

@@ -53,7 +53,7 @@ def test_point_prompt_generator(self):
             generator = PointAndBoxPromptGenerator(n_pos, n_neg, dilation_strength=4)
             for label_id in label_ids:
                 center, box = centers.get(label_id), boxes.get(label_id)
-                coords, point_labels, _, _ = generator(labels, label_id, center, box)
+                coords, point_labels, _, _ = generator(labels, label_id, box, center)
                 coords_ = (np.array([int(coo[0]) for coo in coords]),
                            np.array([int(coo[1]) for coo in coords]))
                 mask = labels == label_id
@@ -80,7 +80,7 @@ def test_box_prompt_generator(self):
 
         for label_id in label_ids:
             center, box_ = centers.get(label_id), boxes.get(label_id)
-            _, _, box, _ = generator(labels, label_id, center, box_)
+            _, _, box, _ = generator(labels, label_id, box_, center)
             coords = np.where(labels == label_id)
             expected_box = [coo.min() for coo in coords] + [coo.max() + 1 for coo in coords]
             self.assertEqual(expected_box, list(box[0]))