Merge pull request #16 from computational-cell-analytics/test-prompts

Add test for the point prompt generator

Author: constantinpape

micro_sam/prompt_generators.py

@@ -32,7 +32,7 @@ def __call__(self, gt, gt_id, center_coordinates, bbox_coordinates):
         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
+        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

micro_sam/util.py

@@ -295,7 +295,7 @@ def compute_iou(mask1, mask2):
     return iou
 
 
-def get_cell_center_coordinates(gt, mode="p"):
+def get_cell_center_coordinates(gt, mode="v"):
     """
     Returns the center coordinates of the foreground instances in the ground-truth
     """
@@ -304,12 +304,14 @@ def get_cell_center_coordinates(gt, mode="p"):
     properties = regionprops(gt)
 
     if mode == "p":
-        center_coordinates = [prop.centroid for prop in properties]
+        center_coordinates = {prop.label: prop.centroid for prop in properties}
     elif mode == "v":
         center_coordinates = vigra.filters.eccentricityCenters(gt.astype('float32'))
+        center_coordinates = {i: coord for i, coord in enumerate(center_coordinates) if i > 0}
 
-    bbox_coordinates = [prop.bbox for prop in properties]
+    bbox_coordinates = {prop.label: prop.bbox for prop in properties}
 
+    assert len(bbox_coordinates) == len(center_coordinates)
     return center_coordinates, bbox_coordinates
 
 

test/test_prompt_generators.py

@@ -0,0 +1,90 @@
+import unittest
+import numpy as np
+
+from skimage.data import binary_blobs
+from skimage.measure import label
+
+
+class TestPromptGenerators(unittest.TestCase):
+
+    def _get_test_data(self):
+        data = binary_blobs(length=256)
+        labels = label(data)
+        return labels
+
+    def _debug(self, mask, center, box, coords, point_labels):
+        import napari
+
+        v = napari.Viewer()
+        v.add_image(mask)
+        v.add_points([center], name="center")
+        v.add_shapes(
+            [np.array(
+                [[box[0], box[1]], [box[2], box[3]]]
+            )],
+            shape_type="rectangle"
+        )
+        prompts = v.add_points(
+            data=np.array(coords),
+            name="prompts",
+            properties={"label": point_labels},
+            edge_color="label",
+            edge_color_cycle=["#00FF00", "#FF0000"],
+            symbol="o",
+            face_color="transparent",
+            edge_width=0.5,
+            size=5,
+            ndim=2
+        )  # this function helps to view the (colored) background/foreground points
+        prompts.edge_color_mode = "cycle"
+        napari.run()
+
+    def test_point_prompt_generator(self):
+        from micro_sam.prompt_generators import PointAndBoxPromptGenerator
+        from micro_sam.util import get_cell_center_coordinates
+
+        labels = self._get_test_data()
+        label_ids = np.unique(labels)[1:]
+
+        centers, boxes = get_cell_center_coordinates(labels)
+
+        test_point_pairs = [(1, 0), (1, 1), (4, 3), (2, 4), (3, 9), (13, 27)]
+        for (n_pos, n_neg) in test_point_pairs:
+            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_ = (np.array([int(coo[0]) for coo in coords]),
+                           np.array([int(coo[1]) for coo in coords]))
+                mask = labels == label_id
+                expected_labels = mask[coords_]
+                agree = (point_labels == expected_labels)
+
+                # DEBUG: check the points in napari if they don't match
+                debug = False
+                if not agree.all() and debug:
+                    print(n_pos, n_neg)
+                    self._debug(mask, center, box, coords, point_labels)
+
+                self.assertTrue(agree.all())
+
+    def test_box_prompt_generator(self):
+        from micro_sam.prompt_generators import PointAndBoxPromptGenerator
+        from micro_sam.util import get_cell_center_coordinates
+
+        labels = self._get_test_data()
+        label_ids = np.unique(labels)[1:]
+
+        centers, boxes = get_cell_center_coordinates(labels)
+        generator = PointAndBoxPromptGenerator(0, 0, dilation_strength=0, get_point_prompts=False, get_box_prompts=True)
+
+        for label_id in label_ids:
+            center, box_ = centers.get(label_id), boxes.get(label_id)
+            _, _, box, _ = generator(labels, label_id, center, box_)
+            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]))
+
+
+if __name__ == "__main__":
+    unittest.main()