Add doc strings and type annotations for prompt_generators

Author: constantinpape

micro_sam/instance_segmentation.py

@@ -1,6 +1,7 @@
 import multiprocessing as mp
 import warnings
-from abc import ABC, Mapping
+from abc import ABC
+from collections.abc import Mapping
 from concurrent import futures
 from copy import deepcopy
 from typing import Any, List, Optional
@@ -1055,8 +1056,6 @@ def segment_tile(_, tile_id):
 # this is still experimental and not yet ready to be integrated within the annotator_3d
 # (will need to see how well it works with retrained models)
 def _segment_instances_from_embeddings_3d(predictor, image_embeddings, verbose=1, iou_threshold=0.50, **kwargs):
-    """
-    """
     if image_embeddings["original_size"] is None:  # tiled embeddings
         is_tiled = True
         image_shape = tuple(image_embeddings["features"].attrs["shape"])

micro_sam/prompt_generators.py

@@ -1,26 +1,76 @@
+from collections.abc import Mapping
+from typing import Optional
+
 import numpy as np
 from scipy.ndimage import binary_dilation
 
 
 class PointAndBoxPromptGenerator:
-    def __init__(self, n_positive_points, n_negative_points, dilation_strength,
-                 get_point_prompts=True, get_box_prompts=False):
+    """Generate point and/or box prompts from an instance segmentation.
+
+    You can use this class to derive prompts from an instance segmentation, either for
+    evaluation purposes or for training Segment Anything on custom data.
+    In order to use this generator you need to precompute the bounding boxes and center
+    coordiantes of the instance segmentation, using e.g. `util.get_bounding_boxes_and_centers`.
+    Here's an example for how to use this class:
+    ```python
+    # Initialize generator for 1 positive and 4 negative point prompts.
+    prompt_generator = PointAndBoxPromptGenerator(1, 4, dilation_strength=8)
+    # Precompute the bounding boxes for the given segmentation
+    bounding_boxes, _ = util.get_bounding_boxes_and_centes(segmentation)
+    # generate point prompts for the object with id 1 in 'segmentation'
+    seg_id = 1
+    points, point_labels, _, _ = prompt_generator(segmentation, seg_id, bounding_boxes)
+    ```
+
+    Args:
+        n_positive_points: The number of positive point prompts to generate per mask.
+        n_negative_points: The number of negative point prompts to generate per mask.
+        dilation_strength: The factor by which the mask is dilated before generating prompts.
+        get_point_prompts: Whether to generate point prompts.
+        get_box_prompts: Whether to generate box prompts.
+    """
+    def __init__(
+        self,
+        n_positive_points: int,
+        n_negative_points: int,
+        dilation_strength: int,
+        get_point_prompts: bool = True,
+        get_box_prompts: bool = 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, bbox_coordinates, center_coordinates=None):
-        """
-        Parameters:
-            gt: True Labels
-            gt_id: Instance ID for the Cells
-            center_coordinates: Coordinates for the centroid seed of the cell
-            bbox_coordinates: Bounding box coordinates around the cell
+            raise ValueError("You need to request box prompts, point prompts or both.")
+
+    def __call__(
+        self,
+        segmentation: np.ndarray,
+        segmentation_id: int,
+        bbox_coordinates: Mapping[int, tuple],
+        center_coordinates: Optional[Mapping[int, np.ndarray]] = None
+    ) -> tuple[
+        Optional[list[tuple]], Optional[list[int]], Optional[list[tuple]], np.ndarray
+    ]:
+        """Generate the prompts for one object in the segmentation.
+
+        Args:
+            segmentation: The instance segmentation.
+            segmentation_id: The ID of the instance.
+            bbox_coordinates: The precomputed bounding boxes of all objects in the segmentation.
+            center_coordinates: The precomputed center coordinates of all objects in the segmentation.
+                If passed, these coordinates will be used as the first positive point prompt.
+                If not passed a random point from within the object mask will be used.
+
+        Returns:
+            List of point coordinates. Returns None, if get_point_prompts is false.
+            List of point labels. Returns None, if get_point_prompts is false.
+            List containing the object bounding box. Returns None, if get_box_prompts is false.
+            Object mask.
         """
         coord_list = []
         label_list = []
@@ -30,7 +80,8 @@ def __call__(self, gt, gt_id, bbox_coordinates, center_coordinates=None):
             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
+            # 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:
@@ -40,7 +91,7 @@ def __call__(self, gt, gt_id, bbox_coordinates, center_coordinates=None):
         if self.get_box_prompts:
             bbox_list = [bbox_coordinates]
 
-        object_mask = gt == gt_id
+        object_mask = segmentation == segmentation_id
 
         if n_positive_remaining > 0:
             # all coordinates of our current object
@@ -50,9 +101,10 @@ def __call__(self, gt, gt_id, bbox_coordinates, center_coordinates=None):
             n_coordinates = len(object_coordinates[0])
 
             # randomly sampling n_positive_remaining_points from these coordinates
-            positive_indices = np.random.choice(n_coordinates, replace=False,
-                                                size=min(n_positive_remaining, n_coordinates)  # handles the cases with insufficient fg pixels
-                                                )
+            positive_indices = np.random.choice(
+                n_coordinates, replace=False,
+                size=min(n_positive_remaining, n_coordinates)  # handles the cases with insufficient fg pixels
+            )
             for positive_index in positive_indices:
                 positive_coordinates = int(object_coordinates[0][positive_index]), \
                     int(object_coordinates[1][positive_index])
@@ -63,10 +115,10 @@ def __call__(self, gt, gt_id, bbox_coordinates, center_coordinates=None):
         # getting the negative points
         # for this we do the opposite and we set the mask to the bounding box - the object mask
         # we need to dilate the object mask before doing this: we use scipy.ndimage.binary_dilation for this
-        dilated_object = binary_dilation(object_mask, iterations=self.dilation_strength)
-        background_mask = np.zeros(gt.shape)
+        dilated_object = binary_dilation(object_mask, iterations=self.dilation_strensegmentationh)
+        background_mask = np.zeros(segmentation.shape)
         background_mask[bbox_coordinates[0]:bbox_coordinates[2], bbox_coordinates[1]:bbox_coordinates[3]] = 1
-        background_mask = binary_dilation(background_mask, iterations=self.dilation_strength)
+        background_mask = binary_dilation(background_mask, iterations=self.dilation_strensegmentationh)
         background_mask = abs(
             background_mask.astype(np.float32) - dilated_object.astype(np.float32)
         )  # casting booleans to do subtraction
@@ -80,9 +132,10 @@ def __call__(self, gt, gt_id, bbox_coordinates, center_coordinates=None):
             n_coordinates = len(background_coordinates[0])
 
             # randomly sample n_positive_remaining_points from these coordinates
-            negative_indices = np.random.choice(n_coordinates, replace=False,
-                                                size=min(n_negative_remaining, n_coordinates)  # handles the cases with insufficient bg pixels
-                                                )
+            negative_indices = np.random.choice(
+                n_coordinates, replace=False,
+                size=min(n_negative_remaining, n_coordinates)  # handles the cases with insufficient bg pixels
+            )
             for negative_index in negative_indices:
                 negative_coordinates = int(background_coordinates[0][negative_index]), \
                     int(background_coordinates[1][negative_index])

micro_sam/util.py

@@ -430,18 +430,28 @@ def compute_iou(mask1, mask2):
     return iou
 
 
-def get_cell_center_coordinates(gt, mode="v"):
-    """
-    Returns the center coordinates of the foreground instances in the ground-truth
+def get_bounding_boxes_and_centers(
+    segmentation: np.ndarray,
+    mode: str = "v"
+) -> tuple[Mapping[int, np.ndarray], Mapping[int, tuple]]:
+    """Returns the center coordinates of the foreground instances in the ground-truth.
+
+    Args:
+        segmentation:
+        mode:
+
+    Returns:
+        A dictionary that maps object ids to the corresponding centroid.
+        A dictionary that maps object_ids to the corresponding bounding box.
     """
     assert mode in ["p", "v"], "Choose either 'p' for regionprops or 'v' for vigra"
 
-    properties = regionprops(gt)
+    properties = regionprops(segmentation)
 
     if mode == "p":
         center_coordinates = {prop.label: prop.centroid for prop in properties}
     elif mode == "v":
-        center_coordinates = vigra.filters.eccentricityCenters(gt.astype('float32'))
+        center_coordinates = vigra.filters.eccentricityCenters(segmentation.astype('float32'))
         center_coordinates = {i: coord for i, coord in enumerate(center_coordinates) if i > 0}
 
     bbox_coordinates = {prop.label: prop.bbox for prop in properties}