Merge branch 'master' into test-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
@@ -74,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

micro_sam/util.py

@@ -105,6 +105,13 @@ def get_sam_model(device=None, model_type="vit_h", checkpoint_path=None, return_
 
 
 def _to_image(input_):
+    # we require the input to be uint8
+    if input_.dtype != np.dtype("uint8"):
+        # first normalize the input to [0, 1]
+        input_ = input_.astype("float32") - input_.min()
+        input_ = input_ / input_.max()
+        # then bring to [0, 255] and cast to uint8
+        input_ = (input_ * 255).astype("uint8")
     if input_.ndim == 2:
         image = np.concatenate([input_[..., None]] * 3, axis=-1)
     elif input_.ndim == 3 and input_.shape[-1] == 3: