Implement segment_from_box_and_points and test

Author: constantinpape

micro_sam/segment_from_prompts.py

@@ -130,3 +130,22 @@ def segment_from_box(
         return mask, scores, logits
     else:
         return mask
+
+
+def segment_from_box_and_points(
+    predictor, box, points, labels,
+    image_embeddings=None, i=None, original_size=None,
+    multimask_output=False, return_all=False,
+):
+    if image_embeddings is not None:
+        util.set_precomputed(predictor, image_embeddings, i)
+    mask, scores, logits = predictor.predict(
+        point_coords=points[:, ::-1],  # SAM has reversed XY conventions
+        point_labels=labels,
+        box=_process_box(box, original_size),
+        multimask_output=multimask_output
+    )
+    if return_all:
+        return mask, scores, logits
+    else:
+        return mask

test/test_segment_from_prompts.py

@@ -7,37 +7,47 @@
 
 
 class TestSegmentFromPrompts(unittest.TestCase):
-    def _get_input(self, shape=(256, 256)):
+    @staticmethod
+    def _get_input(shape=(256, 256)):
         mask = np.zeros(shape, dtype="uint8")
         center = tuple(sh // 2 for sh in shape)
         circle = disk(center, radius=20, shape=shape)
         mask[circle] = 1
         image = mask * 255
         return mask, image
 
-    def _get_model(self, image):
+    @staticmethod
+    def _get_model(image):
         predictor = util.get_sam_model(model_type="vit_b")
         image_embeddings = util.precompute_image_embeddings(predictor, image)
         util.set_precomputed(predictor, image_embeddings)
         return predictor
 
+    # we compute the default mask and predictor once for the class
+    # so that we don't have to precompute it every time
+    @classmethod
+    def setUpClass(cls):
+        cls.mask, cls.image = cls._get_input()
+        cls.predictor = cls._get_model(cls.image)
+
     def test_segment_from_points(self):
         from micro_sam.segment_from_prompts import segment_from_points
 
-        mask, image = self._get_input()
-        predictor = self._get_model(image)
-
         points = np.array([[128, 128], [64, 64], [192, 192], [64, 192], [192, 64]])
         labels = np.array([1, 0, 0, 0, 0])
 
-        predicted = segment_from_points(predictor, points, labels)
-        self.assertGreater(util.compute_iou(mask, predicted), 0.9)
+        predicted = segment_from_points(self.predictor, points, labels)
+        self.assertGreater(util.compute_iou(self.mask, predicted), 0.9)
 
     def _test_segment_from_mask(self, shape=(256, 256), use_mask=True):
         from micro_sam.segment_from_prompts import segment_from_mask
 
-        mask, image = self._get_input(shape)
-        predictor = self._get_model(image)
+        if shape == (256, 256):
+            mask, image = self.mask, self.image
+            predictor = self.predictor
+        else:
+            mask, image = self._get_input(shape)
+            predictor = self._get_model(image)
 
         # with mask and bounding box (default setting)
         if use_mask:
@@ -64,12 +74,19 @@ def test_segment_from_mask_non_square(self):
     def test_segment_from_box(self):
         from micro_sam.segment_from_prompts import segment_from_box
 
-        mask, image = self._get_input()
-        predictor = self._get_model(image)
+        box = np.array([106, 106, 150, 150])
+        predicted = segment_from_box(self.predictor, box)
+        self.assertGreater(util.compute_iou(self.mask, predicted), 0.9)
+
+    def test_segment_from_box_and_points(self):
+        from micro_sam.segment_from_prompts import segment_from_box_and_points
 
         box = np.array([106, 106, 150, 150])
-        predicted = segment_from_box(predictor, box)
-        self.assertGreater(util.compute_iou(mask, predicted), 0.9)
+        points = np.array([[128, 128], [64, 64], [192, 192], [64, 192], [192, 64]])
+        labels = np.array([1, 0, 0, 0, 0])
+
+        predicted = segment_from_box_and_points(self.predictor, box, points, labels)
+        self.assertGreater(util.compute_iou(self.mask, predicted), 0.9)
 
 
 if __name__ == "__main__":