Update segmentation functionality and add more inference scripts

Author: constantinpape

scripts/cooper/full_reconstruction/assort_az_and_vesicles.py

@@ -1,13 +1,36 @@
 import os
 from glob import glob
 
+import h5py
+from tqdm import tqdm
 
-INPUT_ROOT = ""
-OUTPUT_ROOT = ""
+
+INPUT_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/04Dataset_for_vesicle_eval/model_segmentation"  # noqa
+OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/04_full_reconstruction"
+
+
+def assort_az_and_vesicles(in_path, out_path):
+    if os.path.exists(out_path):
+        return
+
+    with h5py.File(in_path, "r") as f:
+        raw = f["raw"][:]
+        vesicles = f["/vesicles/segment_from_combined_vesicles"][:]
+        az = f["/AZ/segment_from_AZmodel_v3"][:]
+
+    os.makedirs(os.path.split(out_path)[0], exist_ok=True)
+    with h5py.File(out_path, "a") as f:
+        f.create_dataset("raw", data=raw, compression="gzip")
+        f.create_dataset("labels/vesicles", data=vesicles, compression="gzip")
+        f.create_dataset("labels/active_zone", data=az, compression="gzip")
 
 
 def main():
-    pass
+    paths = sorted(glob(os.path.join(INPUT_ROOT, "**/*.h5"), recursive=True))
+    for path in tqdm(paths):
+        fname = os.path.relpath(path, INPUT_ROOT)
+        out_path = os.path.join(OUTPUT_ROOT, fname)
+        assort_az_and_vesicles(path, out_path)
 
 
 if __name__ == "__main__":

scripts/cooper/full_reconstruction/segment_compartments.py

@@ -0,0 +1,34 @@
+import os
+from glob import glob
+
+import h5py
+from synaptic_reconstruction.inference.compartments import segment_compartments
+from tqdm import tqdm
+
+ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/04_full_reconstruction"  # noqa
+MODEL_PATH = "/user/pape41/u12086/Work/my_projects/synaptic-reconstruction/scripts/cooper/training/checkpoints/compartment_model_3d/v2"  # noqa
+
+
+def label_transform_3d():
+    pass
+
+
+def run_seg(path):
+    with h5py.File(path, "r") as f:
+        if "labels/compartments" in f:
+            return
+        raw = f["raw"][:]
+
+    scale = (0.25, 0.25, 0.25)
+    seg = segment_compartments(raw, model_path=MODEL_PATH, scale=scale, verbose=False)
+    with h5py.File(path, "a") as f:
+        f.create_dataset("labels/compartments", data=seg, compression="gzip")
+
+
+def main():
+    paths = sorted(glob(os.path.join(ROOT, "**/*.h5"), recursive=True))
+    for path in tqdm(paths):
+        run_seg(path)
+
+
+main()

scripts/cooper/full_reconstruction/segment_mitochondria.py

@@ -0,0 +1,30 @@
+import os
+from glob import glob
+
+import h5py
+from synaptic_reconstruction.inference.mitochondria import segment_mitochondria
+from tqdm import tqdm
+
+ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/04_full_reconstruction"  # noqa
+MODEL_PATH = "/scratch-grete/projects/nim00007/models/exports_for_cooper/mito_model_s2.pt"  # noqa
+
+
+def run_seg(path):
+    with h5py.File(path, "r") as f:
+        if "labels/mitochondria" in f:
+            return
+        raw = f["raw"][:]
+
+    scale = (0.5, 0.5, 0.5)
+    seg = segment_mitochondria(raw, model_path=MODEL_PATH, scale=scale, verbose=False)
+    with h5py.File(path, "a") as f:
+        f.create_dataset("labels/mitochondria", data=seg, compression="gzip")
+
+
+def main():
+    paths = sorted(glob(os.path.join(ROOT, "**/*.h5"), recursive=True))
+    for path in tqdm(paths):
+        run_seg(path)
+
+
+main()

scripts/cooper/full_reconstruction/visualize_results.py

@@ -13,7 +13,7 @@
 from skimage.segmentation import watershed
 from skimage.morphology import remove_small_holes
 
-from synaptic_reconstruction.inference.util import get_prediction, _Scaler
+from synaptic_reconstruction.inference.util import get_prediction, _Scaler, _postprocess_seg_3d
 
 
 def _segment_compartments_2d(
@@ -112,27 +112,6 @@ 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,

synaptic_reconstruction/inference/compartments.py

@@ -5,7 +5,7 @@
 import numpy as np
 import torch
 
-from synaptic_reconstruction.inference.util import apply_size_filter, get_prediction, _Scaler
+from synaptic_reconstruction.inference.util import apply_size_filter, get_prediction, _Scaler, _postprocess_seg_3d
 
 
 def _run_segmentation(
@@ -37,6 +37,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)
     return seg
 
 

synaptic_reconstruction/inference/mitochondria.py

@@ -17,6 +17,9 @@
 import xarray
 
 from elf.io import open_file
+from scipy.ndimage import binary_closing
+from skimage.measure import regionprops
+from skimage.morphology import remove_small_holes
 from skimage.transform import rescale, resize
 from torch_em.util.prediction import predict_with_halo
 from tqdm import tqdm
@@ -465,3 +468,24 @@ def apply_size_filter(
     if verbose:
         print("Size filter in", time.time() - t0, "s")
     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