Implement test for non-square segmentation form mask (currently failing)

Author: constantinpape

test/test_segment_from_prompts.py

@@ -7,10 +7,10 @@
 
 
 class TestSegmentFromPrompts(unittest.TestCase):
-    def _get_input(self):
-        shape = (256, 256)
+    def _get_input(self, shape=(256, 256)):
         mask = np.zeros(shape, dtype="uint8")
-        circle = disk((128, 128), radius=20, shape=shape)
+        center = tuple(sh // 2 for sh in shape)
+        circle = disk(center, radius=20, shape=shape)
         mask[circle] = 1
         image = mask * 255
         return mask, image
@@ -33,24 +33,32 @@ def test_segment_from_points(self):
         predicted = segment_from_points(predictor, points, labels)
         self.assertGreater(util.compute_iou(mask, predicted), 0.9)
 
-    def test_segment_from_mask(self):
+    def _test_segment_from_mask(self, shape=(256, 256)):
         from micro_sam.segment_from_prompts import segment_from_mask
 
-        mask, image = self._get_input()
+        mask, image = self._get_input(shape)
         predictor = self._get_model(image)
 
         # with mask and bounding box (default setting)
         predicted = segment_from_mask(predictor, mask)
         self.assertGreater(util.compute_iou(mask, predicted), 0.9)
 
-        # with  bounding box (default setting)
+        # with bounding box
         predicted = segment_from_mask(predictor, mask, use_mask=False, use_box=True)
         self.assertGreater(util.compute_iou(mask, predicted), 0.9)
 
-        # with  bounding box (default setting)
+        # with mask
         predicted = segment_from_mask(predictor, mask, use_mask=True, use_box=False)
         self.assertGreater(util.compute_iou(mask, predicted), 0.9)
 
+    def test_segment_from_mask(self):
+        self._test_segment_from_mask()
+
+    # FIXME this fails due to shape mismatch in the masks
+    @unittest.expectedFailure
+    def test_segment_from_mask_non_square(self):
+        self._test_segment_from_mask((256, 384))
+
     def test_segment_from_box(self):
         from micro_sam.segment_from_prompts import segment_from_box