Fix tests

Author: constantinpape

development/apg_example.py

@@ -108,7 +108,7 @@ def example_script_wsi():
         # WITH MASK: 34 seconds
         generator = TiledAutomaticPromptGenerator(predictor, decoder)
         generator.initialize(
-            data, image_embeddings=image_embeddings, tile_shape=tile_shape, halo=halo, verbose=True, batch_size=24
+            data, image_embeddings=image_embeddings, tile_shape=tile_shape, halo=halo, verbose=True, batch_size=12
         )
 
         # Processing time: 21:12 min
@@ -177,7 +177,7 @@ def debug_wsi():
 def main():
     # example_script()
     # example_script_tiled()
-    example_script_wsi()
+    # example_script_wsi()
     example_script_3d()
     # debug_wsi()
 

micro_sam/multi_dimensional_segmentation.py

@@ -349,6 +349,8 @@ def merge_instance_segmentation_3d(
 
     # Extract the overlap between slices.
     edges = track_utils.compute_edges_from_overlap(slice_segmentation, verbose=False)
+    if len(edges) == 0:  # Nothing to merge.
+        return slice_segmentation
 
     uv_ids = np.array([[edge["source"], edge["target"]] for edge in edges])
     overlaps = np.array([edge["score"] for edge in edges])

micro_sam/util.py

@@ -841,6 +841,7 @@ def _compute_tiled_features_3d(predictor, input_, tile_shape, halo, f, pbar_init
     msg = "Compute Image Embeddings 3D tiled"
     if mask is None:
         n_tiles_total = n_slices * n_tiles_per_plane
+        tiles_in_mask_per_slice = None
     else:
         tiles_in_mask_per_slice = {}
         for z in range(n_slices):
@@ -867,7 +868,7 @@ def _compute_tiled_features_3d(predictor, input_, tile_shape, halo, f, pbar_init
         pbar_update(len(tile_ids))
 
     if mask is not None:
-        features.attrs["tiles_in_mask"] = tiles_in_mask_per_slice
+        features.attrs["tiles_in_mask"] = {str(z): per_slice for z, per_slice in tiles_in_mask_per_slice.items()}
 
     _write_embedding_signature(f, input_, predictor, tile_shape, halo, input_size=None, original_size=None)
     return features

test/test_automatic_segmentation.py

@@ -1,11 +1,14 @@
 import unittest
 
 import numpy as np
+import torch
 from skimage.draw import disk
 from skimage.measure import label as connected_components
 
 import micro_sam.util as util
 
+HAVE_CUDA = torch.cuda.is_available()
+
 
 class TestAutomaticSegmentation(unittest.TestCase):
     model_type = "vit_t" if util.VIT_T_SUPPORT else "vit_b"
@@ -116,7 +119,7 @@ def test_tiled_instance_segmentation_with_decoder_2d(self):
         )
         self.assertEqual(mask.shape, instances.shape)
 
-    @unittest.skip("Skipping long running tests by default.")
+    @unittest.skipUnless(HAVE_CUDA, "Skipping long running tests unless we have a GPU.")
     def test_automatic_mask_generator_3d(self):
         from micro_sam.automatic_segmentation import automatic_instance_segmentation, get_predictor_and_segmenter
 
@@ -129,7 +132,7 @@ def test_automatic_mask_generator_3d(self):
         )
         self.assertEqual(labels.shape, instances.shape)
 
-    @unittest.skip("Skipping long running tests by default.")
+    @unittest.skipUnless(HAVE_CUDA, "Skipping long running tests unless we have a GPU.")
     def test_tiled_automatic_mask_generator_3d(self):
         from micro_sam.automatic_segmentation import automatic_instance_segmentation, get_predictor_and_segmenter