Minor fix to automatic segmentation in 3d (#882)

Fix to automatic segmentation in 3d for data with channels

Author: anwai98

micro_sam/automatic_segmentation.py

@@ -101,6 +101,13 @@ def automatic_instance_segmentation(
     Returns:
         The segmentation result.
     """
+    # Avoid overwriting already stored segmentations.
+    if output_path is not None:
+        output_path = Path(output_path).with_suffix(".tif")
+        if os.path.exists(output_path):
+            print(f"The segmentation results are already stored at '{os.path.abspath(output_path)}'.")
+            return
+
     # Load the input image file.
     if isinstance(input_path, np.ndarray):
         image_data = input_path
@@ -189,8 +196,8 @@ def automatic_instance_segmentation(
 
     # Save the instance segmentation, if 'output_path' provided.
     if output_path is not None:
-        output_path = Path(output_path).with_suffix(".tif")
         imageio.imwrite(output_path, instances, compression="zlib")
+        print(f"The segmentation results are stored at '{os.path.abspath(output_path)}'.")
 
     if return_embeddings:
         return instances, image_embeddings
@@ -256,6 +263,9 @@ def main():
         help="The device to use for the predictor. Can be one of 'cuda', 'cpu' or 'mps' (only MAC)."
         "By default the most performant available device will be selected."
     )
+    parser.add_argument(
+        "-v", "--verbose", action="store_true", help="Whether to allow verbosity of outputs."
+    )
 
     args, parameter_args = parser.parse_known_args()
 
@@ -314,6 +324,7 @@ def _convert_argval(value):
         tile_shape=args.tile_shape,
         halo=args.halo,
         annotate=args.annotate,
+        verbose=args.verbose,
         **generate_kwargs,
     )
 

micro_sam/instance_segmentation.py

@@ -26,9 +26,6 @@
 
 from nifty.tools import blocking
 
-import elf.parallel as parallel
-from elf.parallel.filters import apply_filter
-
 import segment_anything.utils.amg as amg_utils
 from segment_anything.predictor import SamPredictor
 

micro_sam/multi_dimensional_segmentation.py

@@ -397,7 +397,7 @@ def automatic_3d_segmentation(
         The segmentation.
     """
     offset = 0
-    segmentation = np.zeros(volume.shape, dtype="uint32")
+    segmentation = np.zeros(volume.shape[:3], dtype="uint32")
 
     min_object_size = kwargs.pop("min_object_size", 0)
     image_embeddings = util.precompute_image_embeddings(
@@ -413,15 +413,22 @@ def automatic_3d_segmentation(
     for i in tqdm(range(segmentation.shape[0]), desc="Segment slices", disable=not verbose):
         segmentor.initialize(volume[i], image_embeddings=image_embeddings, verbose=False, i=i)
         seg = segmentor.generate(**kwargs)
-        if len(seg) == 0:
+
+        if isinstance(seg, list) and len(seg) == 0:
             continue
         else:
-            seg = mask_data_to_segmentation(seg, with_background=with_background, min_object_size=min_object_size)
+            if isinstance(seg, list):
+                seg = mask_data_to_segmentation(
+                    seg, with_background=with_background, min_object_size=min_object_size
+                )
+
+            # Set offset for instance per slice.
             max_z = seg.max()
             if max_z == 0:
                 continue
             seg[seg != 0] += offset
             offset = max_z + offset
+
         segmentation[i] = seg
 
     segmentation = merge_instance_segmentation_3d(