Merge pull request #46 from computational-cell-analytics/update-comp-seg

Update compartment segmentation and add code for full vesicle reconstruction

Author: constantinpape

scripts/cooper/full_reconstruction/assort_az_and_vesicles.py

@@ -0,0 +1,37 @@
+import os
+from glob import glob
+
+import h5py
+from tqdm import tqdm
+
+
+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():
+    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__":
+    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

@@ -0,0 +1,123 @@
+import os
+from glob import glob
+
+import h5py
+import napari
+import numpy as np
+import pandas as pd
+
+ROOT = "./04_full_reconstruction"
+TABLE = "/home/pape/Desktop/sfb1286/mboc_synapse/draft_figures/full_reconstruction.xlsx"
+
+# Skip datasets for which all figures were already done.
+SKIP_DS = ["20241019_Tomo-eval_MF_Synapse"]
+
+
+def _get_name_and_row(path, table):
+    ds_name, name = os.path.split(path)
+    ds_name = os.path.split(ds_name)[1]
+    row = table[(table["dataset"] == ds_name) & (table["tomogram"] == name)]
+    return ds_name, name, row
+
+
+def _get_compartment_ids(row):
+    compartment_ids = []
+    for comp in ("Compartment 1", "Compartment 2", "Compartment 3", "Compartment 4"):
+        comp_ids = row[comp].values[0]
+        try:
+            comp_ids = list(map(int, comp_ids.split(", ")))
+        except AttributeError:
+            pass
+
+        if np.isnan(comp_ids).all():
+            compartment_ids.append(None)
+            continue
+
+        if isinstance(comp_ids, int):
+            comp_ids = [comp_ids]
+        compartment_ids.append(comp_ids)
+
+    return compartment_ids
+
+
+def visualize_result(path, table):
+    ds_name, name, row = _get_name_and_row(path, table)
+
+    if ds_name in SKIP_DS:
+        return
+
+    # if row["Use for vis"].values[0] == "yes":
+    if row["Use for vis"].values[0] in ("yes", "no"):
+        return
+    compartment_ids = _get_compartment_ids(row)
+
+    # access = np.s_[:]
+    access = np.s_[::2, ::2, ::2]
+
+    with h5py.File(path, "r") as f:
+        raw = f["raw"][access]
+        vesicles = f["labels/vesicles"][access]
+        active_zone = f["labels/active_zone"][access]
+        mitos = f["labels/mitochondria"][access]
+        compartments = f["labels/compartments"][access]
+
+    if any(comp_ids is not None for comp_ids in compartment_ids):
+        mask = np.zeros(raw.shape, dtype="bool")
+        compartments_new = np.zeros_like(compartments)
+
+        print("Filtering compartments:")
+        for i, comp_ids in enumerate(compartment_ids, 1):
+            if comp_ids is None:
+                continue
+            print(i, comp_ids)
+            this_mask = np.isin(compartments, comp_ids)
+            mask[this_mask] = 1
+            compartments_new[this_mask] = i
+
+        vesicles[~mask] = 0
+        mitos[~mask] = 0
+        compartments = compartments_new
+
+    v = napari.Viewer()
+    v.add_image(raw)
+    v.add_labels(mitos)
+    v.add_labels(vesicles)
+    v.add_labels(compartments)
+    v.add_labels(active_zone)
+    v.title = f"{ds_name}/{name}"
+    napari.run()
+
+
+def visualize_only_compartment(path, table):
+    ds_name, name, row = _get_name_and_row(path, table)
+    compartment_ids = _get_compartment_ids(row)
+
+    # Skip if we already have annotated the presynapse compartment(s)
+    if any(comp_id is not None for comp_id in compartment_ids):
+        print("Compartments already annotated for", ds_name, name)
+        return
+
+    # access = np.s_[:]
+    access = np.s_[::2, ::2, ::2]
+
+    with h5py.File(path, "r") as f:
+        raw = f["raw"][access]
+        compartments = f["labels/compartments"][access]
+
+    v = napari.Viewer()
+    v.add_image(raw)
+    v.add_labels(compartments)
+    v.title = f"{ds_name}/{name}"
+    napari.run()
+
+
+def main():
+    paths = sorted(glob(os.path.join(ROOT, "**/*.h5"), recursive=True))
+    table = pd.read_excel(TABLE)
+    for path in paths:
+        visualize_result(path, table)
+        # visualize_only_compartment(path, table)
+
+
+if __name__ == "__main__":
+    main()

scripts/cooper/ground_truth/compartments/check_predictions.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/ground_truth/compartments/run_annotation_04.py

@@ -0,0 +1,55 @@
+import os
+from glob import glob
+
+import h5py
+import napari
+from magicgui import magicgui
+
+
+def run_annotation(input_path, output_path):
+    with h5py.File(input_path, "r") as f:
+        raw = f["raw"][:]
+        seg = f["labels/compartments"][:]
+
+    v = napari.Viewer()
+    v.add_image(raw)
+    v.add_labels(seg)
+
+    @magicgui(call_button="Save Annotations")
+    def save_annotations():
+        seg = v.layers["seg"].data
+
+        if os.path.exists(output_path):
+            with h5py.File(output_path, "a") as f:
+                f["labels/compartments"][:] = seg
+        else:
+            with h5py.File(output_path, "a") as f:
+                f.create_dataset("raw", data=raw, compression="gzip")
+                f.create_dataset("labels/compartments", data=seg, compression="gzip")
+
+    v.window.add_dock_widget(save_annotations)
+
+    napari.run()
+
+
+def main():
+    inputs = sorted(glob("./predictions/**/*.h5", recursive=True))
+
+    output_folder = "./annotations"
+
+    for input_path in inputs:
+        ds_name, fname = os.path.split(input_path)
+        ds_name = os.path.split(ds_name)[1]
+        ds_folder = os.path.join(output_folder, ds_name)
+        output_path = os.path.join(ds_folder, fname)
+
+        if os.path.exists(output_path):
+            print("Skipping annotations for", output_path)
+            continue
+
+        os.makedirs(ds_folder, exist_ok=True)
+        run_annotation(input_path, output_path)
+
+
+if __name__ == "__main__":
+    main()

scripts/cooper/ground_truth/compartments/run_prediction_04.py

@@ -0,0 +1,64 @@
+import os
+from glob import glob
+
+import h5py
+from tqdm import tqdm
+
+from synaptic_reconstruction.inference.util import _Scaler
+from synaptic_reconstruction.inference.compartments import segment_compartments
+
+INPUT_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/ground_truth/04Dataset_for_vesicle_eval"  # noqa
+# MODEL_PATH = "/mnt/lustre-emmy-hdd/projects/nim00007/compartment_models/compartment_model_3d.pt"
+MODEL_PATH = "/user/pape41/u12086/Work/my_projects/synaptic-reconstruction/scripts/cooper/training/checkpoints/compartment_model_3d/v2"  # noqa
+OUTPUT = "/mnt/lustre-emmy-hdd/projects/nim00007/compartment_predictions"
+
+
+def label_transform_3d():
+    pass
+
+
+def segment_volume(input_path, model_path):
+    with h5py.File(input_path, "r") as f:
+        raw = f["raw"][:]
+
+    scale = (0.25, 0.25, 0.25)
+    scaler = _Scaler(scale, verbose=False)
+    raw = scaler.scale_input(raw)
+
+    n_slices_exclude = 2
+    seg, pred = segment_compartments(
+        raw, model_path, verbose=False, n_slices_exclude=n_slices_exclude, return_predictions=True
+    )
+    # raw, seg = raw[n_slices_exclude:-n_slices_exclude], seg[n_slices_exclude:-n_slices_exclude]
+
+    return raw, seg, pred
+
+
+def main():
+    inputs = sorted(glob(os.path.join(INPUT_ROOT, "**/*.h5"), recursive=True))
+    inputs = [inp for inp in inputs if "cropped_for_2D" not in inp]
+
+    for input_path in tqdm(inputs, desc="Run prediction for 04."):
+        ds_name, fname = os.path.split(input_path)
+        ds_name = os.path.split(ds_name)[1]
+        output_folder = os.path.join(OUTPUT, "segmentation", ds_name)
+        output_path = os.path.join(output_folder, fname)
+
+        if os.path.exists(output_path):
+            continue
+
+        pred_folder = os.path.join(OUTPUT, "prediction", ds_name)
+        os.makedirs(pred_folder, exist_ok=True)
+        pred_path = os.path.join(pred_folder, fname)
+
+        raw, seg, pred = segment_volume(input_path, MODEL_PATH)
+        os.makedirs(output_folder, exist_ok=True)
+        with h5py.File(output_path, "a") as f:
+            f.create_dataset("raw", data=raw, compression="gzip")
+            f.create_dataset("labels/compartments", data=seg, compression="gzip")
+        with h5py.File(pred_path, "a") as f:
+            f.create_dataset("prediction", data=pred, compression="gzip")
+
+
+if __name__ == "__main__":
+    main()