Fix tests if vit_t not available

Author: constantinpape

test/test_instance_segmentation.py

@@ -10,6 +10,7 @@
 
 
 class TestInstanceSegmentation(unittest.TestCase):
+    model_type = "vit_t" if util.VIT_T_SUPPORT else "vit_b"
     embedding_path = "./tmp_embeddings.zarr"
     tile_shape = (512, 512)
     halo = (96, 96)
@@ -37,8 +38,8 @@ def write_object(center, radius):
         return mask, image
 
     @staticmethod
-    def _get_model(image):
-        predictor = util.get_sam_model(model_type="vit_t")
+    def _get_model(image, model_type):
+        predictor = util.get_sam_model(model_type=model_type)
         image_embeddings = util.precompute_image_embeddings(predictor, image)
         return predictor, image_embeddings
 
@@ -47,7 +48,7 @@ def _get_model(image):
     @classmethod
     def setUpClass(cls):
         cls.mask, cls.image = cls._get_input()
-        cls.predictor, cls.image_embeddings = cls._get_model(cls.image)
+        cls.predictor, cls.image_embeddings = cls._get_model(cls.image, cls.model_type)
         cls.large_mask, cls.large_image = cls._get_input(shape=(1024, 1024))
         cls.tiled_embeddings = util.precompute_image_embeddings(
             cls.predictor, cls.large_image, save_path=cls.embedding_path, tile_shape=cls.tile_shape, halo=cls.halo

test/test_prompt_based_segmentation.py

@@ -7,6 +7,8 @@
 
 
 class TestPromptBasedSegmentation(unittest.TestCase):
+    model_type = "vit_t" if util.VIT_T_SUPPORT else "vit_b"
+
     @staticmethod
     def _get_input(shape=(256, 256)):
         mask = np.zeros(shape, dtype="uint8")
@@ -17,8 +19,8 @@ def _get_input(shape=(256, 256)):
         return mask, image
 
     @staticmethod
-    def _get_model(image):
-        predictor = util.get_sam_model(model_type="vit_t")
+    def _get_model(image, model_type):
+        predictor = util.get_sam_model(model_type=model_type)
         image_embeddings = util.precompute_image_embeddings(predictor, image)
         util.set_precomputed(predictor, image_embeddings)
         return predictor
@@ -28,7 +30,7 @@ def _get_model(image):
     @classmethod
     def setUpClass(cls):
         cls.mask, cls.image = cls._get_input()
-        cls.predictor = cls._get_model(cls.image)
+        cls.predictor = cls._get_model(cls.image, cls.model_type)
 
     def test_segment_from_points(self):
         from micro_sam.prompt_based_segmentation import segment_from_points
@@ -59,7 +61,7 @@ def _test_segment_from_mask(self, shape=(256, 256)):
             expected_iou_mask = 0.9
         else:
             mask, image = self._get_input(shape)
-            predictor = self._get_model(image)
+            predictor = self._get_model(image, self.model_type)
             expected_iou_mask = 0.8
 
         #

test/test_training.py

@@ -9,8 +9,10 @@
 import torch_em
 
 from micro_sam.sample_data import synthetic_data
+from micro_sam.util import VIT_T_SUPPORT
 
 
+@unittest.skipUnless(VIT_T_SUPPORT, "Integration test is only run with vit_t support, otherwise it takes too long.")
 class TestTraining(unittest.TestCase):
     """Integration test for training a SAM model.
     """

test/test_util.py

@@ -34,9 +34,9 @@ def test_compute_iou(self):
             self.assertTrue(0.0 < compute_iou(x1, x2) < 1.0)
 
     def test_tiled_prediction(self):
-        from micro_sam.util import precompute_image_embeddings, get_sam_model
+        from micro_sam.util import precompute_image_embeddings, get_sam_model, VIT_T_SUPPORT
 
-        predictor = get_sam_model(model_type="vit_t")
+        predictor = get_sam_model(model_type="vit_t" if VIT_T_SUPPORT else "vit_b")
 
         tile_shape, halo = (256, 256), (16, 16)
         input_ = np.random.rand(512, 512).astype("float32")