added boundary threshold and area threshold to mitochondrai segmentation (#94)

Update postprocessing for mito segmentation

Author: lufre1

synapse_net/inference/mitochondria.py

@@ -13,10 +13,11 @@ def _run_segmentation(
     # blocking shapes for parallel computation
     block_shape=(128, 256, 256),
     halo=(48, 48, 48),
-    seed_distance=6
+    seed_distance=6,
+    boundary_threshold=0.25,
+    area_threshold=5000,
 ):
     t0 = time.time()
-    boundary_threshold = 0.25
     dist = parallel.distance_transform(
         boundaries < boundary_threshold, halo=halo, verbose=verbose, block_shape=block_shape
     )
@@ -43,7 +44,7 @@ def _run_segmentation(
         print("Compute watershed in", time.time() - t0, "s")
 
     seg = apply_size_filter(seg, min_size, verbose=verbose, block_shape=block_shape)
-    seg = _postprocess_seg_3d(seg, area_threshold=5000)
+    seg = _postprocess_seg_3d(seg, area_threshold=area_threshold)
     return seg
 
 
@@ -61,6 +62,8 @@ def segment_mitochondria(
     seed_distance: int = 6,
     ws_block_shape: Tuple[int, ...] = (128, 256, 256),
     ws_halo: Tuple[int, ...] = (48, 48, 48),
+    boundary_threshold: float = 0.25,
+    area_threshold: int = 5000,
 ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
     """Segment mitochondria in an input volume.
 
@@ -78,6 +81,9 @@ def segment_mitochondria(
         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.
+        boundary_threshold: The boundary threshold distance calculation.
+        area_threshold: The maximum area (in pixels) of holes to be removed or filled in the segmentation.
+            This parameter is passed to `skimage.morphology.remove_small_holes`.
 
     Returns:
         The segmentation mask as a numpy array, or a tuple containing the segmentation mask
@@ -97,7 +103,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, seed_distance=seed_distance,
-                            block_shape=ws_block_shape, halo=ws_halo)
+                            block_shape=ws_block_shape, halo=ws_halo, boundary_threshold=boundary_threshold,
+                            area_threshold=area_threshold)
     seg = scaler.rescale_output(seg, is_segmentation=True)
 
     if return_predictions:

synapse_net/inference/util.py

@@ -527,7 +527,7 @@ def _postprocess_seg_3d(seg, area_threshold=1000, iterations=4, iterations_3d=8)
     props = regionprops(seg)
     for prop in props:
         # Get bounding box and mask.
-        bb = tuple(slice(start, stop) for start, stop in zip(prop.bbox[:2], prop.bbox[2:]))
+        bb = tuple(slice(start, stop) for start, stop in zip(prop.bbox[:3], prop.bbox[3:]))
         mask = seg[bb] == prop.label
 
         # Fill small holes and apply closing.