added seed_distance to parameters of mitochondria segmentation (#93)

Add seed_distance and block_shape to parameters of mitochondria segmentation

---------

Co-authored-by: Constantin Pape <[email protected]>

environment.yaml

@@ -15,5 +15,6 @@ dependencies:
     - torch_em
     - torchvision
     - trimesh
+    - zarr < 3
     - pip:
         - napari-skimage-regionprops

synapse_net/inference/mitochondria.py

@@ -12,7 +12,8 @@ def _run_segmentation(
     foreground, boundaries, verbose, min_size,
     # blocking shapes for parallel computation
     block_shape=(128, 256, 256),
-    halo=(48, 48, 48)
+    halo=(48, 48, 48),
+    seed_distance=6
 ):
     t0 = time.time()
     boundary_threshold = 0.25
@@ -24,19 +25,11 @@ def _run_segmentation(
 
     # Get the segmentation via seeded watershed.
     t0 = time.time()
-    seed_distance = 6
     seeds = np.logical_and(foreground > 0.5, dist > seed_distance)
     seeds = parallel.label(seeds, block_shape=block_shape, verbose=verbose)
     if verbose:
         print("Compute connected components in", time.time() - t0, "s")
 
-    # import napari
-    # v = napari.Viewer()
-    # v.add_image(boundaries)
-    # v.add_image(dist)
-    # v.add_labels(seeds)
-    # napari.run()
-
     t0 = time.time()
     hmap = boundaries + ((dist.max() - dist) / dist.max())
     mask = (foreground + boundaries) > 0.5
@@ -65,6 +58,9 @@ def segment_mitochondria(
     return_predictions: bool = False,
     scale: Optional[List[float]] = None,
     mask: Optional[np.ndarray] = None,
+    seed_distance: int = 6,
+    ws_block_shape: Tuple[int, ...] = (128, 256, 256),
+    ws_halo: Tuple[int, ...] = (48, 48, 48),
 ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
     """Segment mitochondria in an input volume.
 
@@ -79,6 +75,9 @@ def segment_mitochondria(
         return_predictions: Whether to return the predictions (foreground, boundaries) alongside the segmentation.
         scale: The scale factor to use for rescaling the input volume before prediction.
         mask: An optional mask that is used to restrict the segmentation.
+        seed_distance: The distance threshold for the seeded watershed.
+        ws_block_shape: The block shape for the seeded watershed.
+        ws_halo: The halo for the seeded watershed.
 
     Returns:
         The segmentation mask as a numpy array, or a tuple containing the segmentation mask
@@ -97,7 +96,8 @@ def segment_mitochondria(
 
     # Run segmentation and rescale the result if necessary.
     foreground, boundaries = pred[:2]
-    seg = _run_segmentation(foreground, boundaries, verbose=verbose, min_size=min_size)
+    seg = _run_segmentation(foreground, boundaries, verbose=verbose, min_size=min_size, seed_distance=seed_distance,
+                            block_shape=ws_block_shape, halo=ws_halo)
     seg = scaler.rescale_output(seg, is_segmentation=True)
 
     if return_predictions: