Update compartment segmentation

Author: constantinpape

scripts/cooper/full_reconstruction/assort_az_and_vesicles.py

@@ -0,0 +1,14 @@
+import os
+from glob import glob
+
+
+INPUT_ROOT = ""
+OUTPUT_ROOT = ""
+
+
+def main():
+    pass
+
+
+if __name__ == "__main__":
+    main()

scripts/cooper/full_reconstruction/segment_compartments.py

@@ -0,0 +1,47 @@
+import os
+from glob import glob
+
+import h5py
+import napari
+
+from synaptic_reconstruction.inference.compartments import _segment_compartments_3d
+
+
+def check_pred(path, pred_path, name):
+    with h5py.File(path, "r") as f:
+        raw = f["raw"][:]
+        # seg = f["labels/compartments"][:]
+
+    with h5py.File(pred_path, "r") as f:
+        pred = f["prediction"][:]
+
+    print("Run segmentation ...")
+    seg_new = _segment_compartments_3d(pred)
+    print("done")
+
+    v = napari.Viewer()
+    v.add_image(raw)
+    v.add_image(pred, visible=False)
+    # v.add_labels(seg, visible=False)
+    v.add_labels(seg_new)
+    v.title = name
+    napari.run()
+
+
+def main():
+    seg_paths = sorted(glob("./predictions/segmentation/**/*.h5", recursive=True))
+
+    for seg_path in seg_paths:
+        ds_name, fname = os.path.split(seg_path)
+        ds_name = os.path.split(ds_name)[1]
+
+        # if ds_name in ("20241019_Tomo-eval_MF_Synapse", "20241019_Tomo-eval_PS_Synapse"):
+        #     continue
+
+        name = f"{ds_name}/{fname}"
+        pred_path = os.path.join("./predictions/prediction", ds_name, fname)
+        assert os.path.exists(pred_path), pred_path
+        check_pred(seg_path, pred_path, name)
+
+
+main()

scripts/cooper/full_reconstruction/segment_mitochondria.py

@@ -8,49 +8,76 @@
 import elf.segmentation as eseg
 import nifty
 from elf.tracking.tracking_utils import compute_edges_from_overlap
-from scipy.ndimage import distance_transform_edt
+from scipy.ndimage import distance_transform_edt, binary_closing
 from skimage.measure import label, regionprops
 from skimage.segmentation import watershed
+from skimage.morphology import remove_small_holes
 
-from synaptic_reconstruction.inference.util import apply_size_filter, get_prediction, _Scaler
-
-
-def _multicut(ws, prediction, beta, n_threads):
-    rag = eseg.features.compute_rag(ws, n_threads=n_threads)
-    edge_features = eseg.features.compute_boundary_mean_and_length(rag, prediction, n_threads=n_threads)
-    edge_probs, edge_sizes = edge_features[:, 0], edge_features[:, 1]
-    edge_costs = eseg.multicut.compute_edge_costs(edge_probs, edge_sizes=edge_sizes, beta=beta)
-    node_labels = eseg.multicut.multicut_kernighan_lin(rag, edge_costs)
-    seg = eseg.features.project_node_labels_to_pixels(rag, node_labels, n_threads)
-    return seg
+from synaptic_reconstruction.inference.util import get_prediction, _Scaler
 
 
 def _segment_compartments_2d(
-    prediction,
-    distances=None,
-    boundary_threshold=0.4,
-    beta=0.6,
-    n_threads=1,
-    run_multicut=True,
-    min_size=500,
+    boundaries,
+    boundary_threshold=0.4,  # Threshold for the boundary distance computation.
+    large_seed_distance=30,  # The distance threshold for computing large seeds (= components).
+    distances=None,  # Pre-computed distances to take into account z-context.
 ):
+    # Compoute distances if already not precomputed.
     if distances is None:
-        distances = distance_transform_edt(prediction < boundary_threshold).astype("float32")
+        distances = distance_transform_edt(boundaries < boundary_threshold).astype("float32")
+        distances_z = distances
+    else:
+        # If the distances were pre-computed then compute them again in 2d.
+        # This is needed for inserting small seeds from maxima, otherwise we will get spurious maxima.
+        distances_z = distance_transform_edt(boundaries < boundary_threshold).astype("float32")
+
+    # Find the large seeds as connected components in the distances > large_seed_distance.
+    seeds = label(distances > large_seed_distance)
+
+    # Remove to small large seeds.
+    min_seed_area = 50
+    ids, sizes = np.unique(seeds, return_counts=True)
+    remove_ids = ids[sizes < min_seed_area]
+    seeds[np.isin(seeds, remove_ids)] = 0
+
+    # Compute the small seeds = local maxima of the in-plane distance map
+    small_seeds = vigra.analysis.localMaxima(distances_z, marker=np.nan, allowAtBorder=True, allowPlateaus=True)
+    small_seeds = label(np.isnan(small_seeds))
+
+    # We only keep small seeds that don't intersect with a large seed.
+    props = regionprops(small_seeds, seeds)
+    keep_seeds = [prop.label for prop in props if prop.max_intensity == 0]
+    keep_mask = np.isin(small_seeds, keep_seeds)
+
+    # Add up the small seeds we keep with the large seeds.
+    all_seeds = seeds.copy()
+    seed_offset = seeds.max()
+    all_seeds[keep_mask] = (small_seeds[keep_mask] + seed_offset)
+
+    # Run watershed to get the segmentation.
+    hmap = boundaries + (distances.max() - distances) / distances.max()
+    raw_segmentation = watershed(hmap, markers=all_seeds)
+
+    # Thee are the large seed ids that we will keep.
+    keep_ids = list(range(1, seed_offset + 1))
+
+    # Iterate over the ids, only keep large seeds and remove holes in their respective masks.
+    props = regionprops(raw_segmentation)
+    segmentation = np.zeros_like(raw_segmentation)
+    for prop in props:
+        if prop.label not in keep_ids:
+            continue
 
-    # replace with skimage?
-    maxima = vigra.analysis.localMaxima(distances, marker=np.nan, allowAtBorder=True, allowPlateaus=True)
-    maxima = label(np.isnan(maxima))
+        # Get bounding box and mask.
+        bb = tuple(slice(start, stop) for start, stop in zip(prop.bbox[:2], prop.bbox[2:]))
+        mask = raw_segmentation[bb] == prop.label
 
-    hmap = distances
-    hmap = (hmap.max() - hmap)
-    hmap /= hmap.max()
-    hmap_ws = hmap + prediction
-    ws = watershed(hmap_ws, markers=maxima)
+        # Fill small holes and apply closing.
+        mask = remove_small_holes(mask, area_threshold=500)
+        mask = np.logical_or(binary_closing(mask, iterations=4), mask)
+        segmentation[bb][mask] = prop.label
 
-    hmap_mc = 0.8 * prediction + 0.2 * hmap
-    seg = _multicut(ws, hmap_mc, beta, n_threads)
-    seg = apply_size_filter(seg, min_size)
-    return seg
+    return segmentation
 
 
 def _merge_segmentation_3d(seg_2d, beta=0.5, min_z_extent=10):
@@ -63,13 +90,12 @@ def _merge_segmentation_3d(seg_2d, beta=0.5, min_z_extent=10):
     graph = nifty.graph.undirectedGraph(n_nodes)
     graph.insertEdges(uv_ids)
 
-    costs = eseg.multicut.compute_edge_costs(overlaps)
+    costs = eseg.multicut.compute_edge_costs(1.0 - overlaps)
     # set background weights to be maximally repulsive
     bg_edges = (uv_ids == 0).any(axis=1)
     costs[bg_edges] = -8.0
 
-    node_labels = eseg.multicut.multicut_decomposition(graph, -1 * costs, beta=beta)
-
+    node_labels = eseg.multicut.multicut_decomposition(graph, costs, beta=beta)
     segmentation = nifty.tools.take(node_labels, seg_2d)
 
     if min_z_extent is not None and min_z_extent > 0:
@@ -86,25 +112,57 @@ def _merge_segmentation_3d(seg_2d, beta=0.5, min_z_extent=10):
     return segmentation
 
 
+def _postprocess_seg_3d(seg):
+    # Structure lement for 2d dilation in 3d.
+    structure_element = np.ones((3, 3))  # 3x3 structure for XY plane
+    structure_3d = np.zeros((1, 3, 3))  # Only applied in the XY plane
+    structure_3d[0] = structure_element
+
+    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:]))
+        mask = seg[bb] == prop.label
+
+        # Fill small holes and apply closing.
+        mask = remove_small_holes(mask, area_threshold=1000)
+        mask = np.logical_or(binary_closing(mask, iterations=4), mask)
+        mask = np.logical_or(binary_closing(mask, iterations=8, structure=structure_3d), mask)
+        seg[bb][mask] = prop.label
+
+    return seg
+
+
 def _segment_compartments_3d(
     prediction,
     boundary_threshold=0.4,
-    n_slices_exclude=5,
+    n_slices_exclude=0,
     min_z_extent=10,
 ):
     distances = distance_transform_edt(prediction < boundary_threshold).astype("float32")
     seg_2d = np.zeros(prediction.shape, dtype="uint32")
 
     offset = 0
+    # Parallelize?
     for z in range(seg_2d.shape[0]):
         if z < n_slices_exclude or z >= seg_2d.shape[0] - n_slices_exclude:
             continue
-        seg_z = _segment_compartments_2d(prediction[z], distances=distances[z], run_multicut=True, min_size=500)
+        seg_z = _segment_compartments_2d(prediction[z], distances=distances[z])
         seg_z[seg_z != 0] += offset
         offset = int(seg_z.max())
         seg_2d[z] = seg_z
 
     seg = _merge_segmentation_3d(seg_2d, min_z_extent)
+    seg = _postprocess_seg_3d(seg)
+
+    # import napari
+    # v = napari.Viewer()
+    # v.add_image(prediction)
+    # v.add_image(distances)
+    # v.add_labels(seg_2d)
+    # v.add_labels(seg)
+    # napari.run()
+
     return seg
 
 
@@ -116,7 +174,7 @@ def segment_compartments(
     verbose: bool = True,
     return_predictions: bool = False,
     scale: Optional[List[float]] = None,
-    mask: Optional[np.ndarray] = None,
+    **kwargs,
 ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
     """
     Segment synaptic compartments in an input volume.
@@ -159,6 +217,7 @@ def segment_compartments(
         seg = _segment_compartments_3d(pred)
     if verbose:
         print("Run segmentation in", time.time() - t0, "s")
+
     seg = scaler.rescale_output(seg, is_segmentation=True)
 
     if return_predictions: