Implement vesicle pool correction WIP

Author: constantinpape

scripts/otoferlin/automatic_processing.py

@@ -2,18 +2,17 @@
 
 import h5py
 import numpy as np
-import pandas as pd
 
 from skimage.measure import label
 from skimage.segmentation import relabel_sequential
 
-from synapse_net.distance_measurements import measure_segmentation_to_object_distances, load_distances
+from synapse_net.distance_measurements import measure_segmentation_to_object_distances
 from synapse_net.file_utils import read_mrc
 from synapse_net.inference.vesicles import segment_vesicles
 from synapse_net.tools.util import get_model, compute_scale_from_voxel_size, _segment_ribbon_AZ
 from tqdm import tqdm
 
-from common import STRUCTURE_NAMES, get_all_tomograms, get_seg_path, get_adapted_model
+from common import get_all_tomograms, get_seg_path, get_adapted_model
 
 # These are tomograms for which the sophisticated membrane processing fails.
 # In this case, we just select the largest boundary piece.
@@ -159,74 +158,6 @@ def postprocess_vesicles(mrc_path, output_path, process_center_crop):
         f.create_dataset(key, data=vesicles, compression="gzip")
 
 
-def measure_distances(mrc_path, seg_path, output_folder):
-    result_folder = os.path.join(output_folder, "distances")
-    if os.path.exists(result_folder):
-        return
-
-    # Get the voxel size.
-    _, voxel_size = read_mrc(mrc_path)
-    resolution = tuple(voxel_size[ax] for ax in "zyx")
-
-    # Load the segmentations.
-    with h5py.File(seg_path, "r") as f:
-        g = f["segmentation"]
-        vesicles = g["vesicles"][:]
-        structures = {name: g[name][:] for name in STRUCTURE_NAMES}
-
-    # Measure all the object distances.
-    os.makedirs(result_folder, exist_ok=True)
-    for name, seg in structures.items():
-        if seg.sum() == 0:
-            print(name, "was not found, skipping the distance computation.")
-            continue
-        print("Compute vesicle distances to", name)
-        save_path = os.path.join(result_folder, f"{name}.npz")
-        measure_segmentation_to_object_distances(vesicles, seg, save_path=save_path, resolution=resolution)
-
-
-def assign_vesicle_pools(output_folder):
-    assignment_path = os.path.join(output_folder, "vesicle_pools.csv")
-    if os.path.exists(assignment_path):
-        return
-
-    distance_folder = os.path.join(output_folder, "distances")
-    distance_paths = {name: os.path.join(distance_folder, f"{name}.npz") for name in STRUCTURE_NAMES}
-    if not all(os.path.exists(path) for path in distance_paths.values()):
-        print("Skip vesicle pool assignment, because some distances are missing.")
-        print("This is probably due to the fact that the corresponding structures were not found.")
-        return
-    distances = {name: load_distances(path) for name, path in distance_paths.items()}
-
-    # The distance criteria.
-    rav_ribbon_distance = 80  # nm
-    mpv_pd_distance = 100  # nm
-    mpv_mem_distance = 50  # nm
-    docked_pd_distance = 100  # nm
-    docked_mem_distance = 2  # nm
-
-    rav_distances, seg_ids = distances["ribbon"][0], np.array(distances["ribbon"][-1])
-    rav_ids = seg_ids[rav_distances < rav_ribbon_distance]
-
-    pd_distances, mem_distances = distances["PD"][0], distances["membrane"][0]
-    assert len(pd_distances) == len(mem_distances) == len(rav_distances)
-
-    mpv_ids = seg_ids[np.logical_and(pd_distances < mpv_pd_distance, mem_distances < mpv_mem_distance)]
-    docked_ids = seg_ids[np.logical_and(pd_distances < docked_pd_distance, mem_distances < docked_mem_distance)]
-
-    # Create a dictionary to map vesicle ids to their corresponding pool.
-    # (RA-V get's over-written by MP-V, which is correct).
-    pool_assignments = {vid: "RA-V" for vid in rav_ids}
-    pool_assignments.update({vid: "MP-V" for vid in mpv_ids})
-    pool_assignments.update({vid: "Docked-V" for vid in docked_ids})
-
-    pool_assignments = pd.DataFrame({
-        "vesicle_id": list(pool_assignments.keys()),
-        "pool": list(pool_assignments.values()),
-    })
-    pool_assignments.to_csv(assignment_path, index=False)
-
-
 def process_tomogram(mrc_path):
     output_path = get_seg_path(mrc_path)
     output_folder = os.path.split(output_path)[0]
@@ -238,13 +169,6 @@ def process_tomogram(mrc_path):
     process_ribbon_structures(mrc_path, output_path, process_center_crop)
     postprocess_vesicles(mrc_path, output_path, process_center_crop)
 
-    # We don't need to do the analysis of the auto semgentation, it only
-    # makes sense to do this after segmentation. I am leaving this here for
-    # now, to move it to the files for analysis later.
-
-    # measure_distances(mrc_path, output_path, output_folder)
-    # assign_vesicle_pools(output_folder)
-
 
 def main():
     tomograms = get_all_tomograms()

scripts/otoferlin/check_automatic_result.py

@@ -3,7 +3,6 @@
 import h5py
 import napari
 import numpy as np
-import pandas as pd
 
 from synapse_net.file_utils import read_mrc
 from skimage.exposure import equalize_adapthist
@@ -12,22 +11,6 @@
 from common import get_all_tomograms, get_seg_path, get_colormaps
 
 
-def _get_vesicle_pools(seg, assignment_path):
-    assignments = pd.read_csv(assignment_path)
-    pool_names = pd.unique(assignments.pool).tolist()
-    pools = np.zeros_like(seg)
-
-    pool_colors = get_colormaps()["pools"]
-    colormap = {}
-    for pool_id, pool_name in enumerate(pool_names, 1):
-        pool_vesicle_ids = assignments[assignments.pool == pool_name].vesicle_id
-        pool_mask = np.isin(seg, pool_vesicle_ids)
-        pools[pool_mask] = pool_id
-        colormap[pool_id] = pool_colors[pool_name]
-
-    return pools, colormap
-
-
 def check_automatic_result(mrc_path, version, use_clahe=False, center_crop=True, segmentation_group="segmentation"):
     tomogram, _ = read_mrc(mrc_path)
     if center_crop:
@@ -53,11 +36,6 @@ def check_automatic_result(mrc_path, version, use_clahe=False, center_crop=True,
                 segmentations[name] = ds[bb]
                 colormaps[name] = get_colormaps().get(name, None)
 
-    output_folder = os.path.split(seg_path)[0]
-    assignment_path = os.path.join(output_folder, "vesicle_pools.csv")
-    if os.path.exists(assignment_path) and "vesicles" in segmentations:
-        segmentations["pools"], colormaps["pools"] = _get_vesicle_pools(segmentations["vesicles"], assignment_path)
-
     v = napari.Viewer()
     v.add_image(tomogram)
     for name, seg in segmentations.items():

scripts/otoferlin/common.py

@@ -1,6 +1,8 @@
 import os
 from glob import glob
 
+import imageio.v3 as imageio
+import h5py
 from synapse_net.tools.util import load_custom_model
 
 
@@ -62,16 +64,29 @@ def get_colormaps():
         "Docked-V": (1, 1, 0),
         None: "gray",
     }
+    ribbon_map = {1: "red", None: "gray"}
     membrane_map = {1: "purple", None: "gray"}
     pd_map = {1: "magenta", None: "gray"}
-    return {"pools": pool_map, "membrane": membrane_map, "PD": pd_map}
-
-
-# TODO: sync the ukon folder with the tomograms.
-# UKON Path:
-# /run/user/1000/gvfs/smb-share:server=wfs-medizin.top.gwdg.de,share=ukon-all$/UKON100/archiv/EM/For Segmentation
-def sync_tomograms():
-    pass
+    return {"pools": pool_map, "membrane": membrane_map, "PD": pd_map, "ribbon": ribbon_map}
+
+
+def load_segmentations(seg_path):
+    # Keep the typo in the name, as these are the hdf5 keys!
+    seg_names = {"vesicles": "veiscles_postprocessed"}
+    seg_names.update({name: name for name in STRUCTURE_NAMES})
+
+    segmentations = {}
+    correction_folder = os.path.join(os.path.split(seg_path)[0], "correction")
+    with h5py.File(seg_path, "r") as f:
+        g = f["segmentation"]
+        for out_name, name in seg_names.items():
+            correction_path = os.path.join(correction_folder, f"{name}.tif")
+            if os.path.exists(correction_path):
+                print("Loading corrected", name, "segmentation from", correction_path)
+                segmentations[out_name] = imageio.imread(correction_path)
+            else:
+                segmentations[out_name] = g[f"{name}"][:]
+    return segmentations
 
 
 if __name__ == "__main__":

scripts/otoferlin/correct_structure_segmentation.py

@@ -1,41 +1,25 @@
-import os
 from pathlib import Path
 
-import imageio.v3 as imageio
-import h5py
 import napari
 
 from synapse_net.file_utils import read_mrc
-from common import get_all_tomograms, get_seg_path
+from common import get_all_tomograms, get_seg_path, load_segmentations, get_colormaps
 
 
 def correct_structure_segmentation(mrc_path):
     seg_path = get_seg_path(mrc_path)
 
     data, _ = read_mrc(mrc_path)
-    correction_folder = os.path.join(os.path.split(seg_path)[0], "correction")
-    fname = Path(mrc_path).stem
-
-    names = ("ribbon", "PD", "membrane", "veiscles_postprocessed")
-    segmentations = {}
-    with h5py.File(seg_path, "r") as f:
-        for name in names:
-            correction_path = os.path.join(correction_folder, f"{name}.tif")
-            if os.path.exists(correction_path):
-                print("Loading segmentation for", name, "from", correction_path)
-                segmentations[name] = imageio.imread(correction_path)
-            else:
-                segmentations[name] = f[f"segmentation/{name}"][:]
-    color_maps = {
-        "ribbon": {1: "red", None: "gray"},
-        "PD": {1: "purple", None: "gray"},
-        "membrane": {1: "magenta", None: "gray"},
-    }
+    segmentations = load_segmentations(seg_path)
+    color_maps = get_colormaps()
 
     v = napari.Viewer()
     v.add_image(data)
     for name, seg in segmentations.items():
+        if name == "vesicles":
+            name = "veiscles_postprocessed"
         v.add_labels(seg, name=name, colormap=color_maps.get(name, None))
+    fname = Path(mrc_path).stem
     v.title = fname
     napari.run()
 

scripts/otoferlin/correct_vesicle_pools.py

@@ -0,0 +1,100 @@
+import os
+
+import imageio.v3 as imageio
+import napari
+import numpy as np
+import pandas as pd
+from magicgui import magicgui
+
+from synapse_net.file_utils import read_mrc
+from skimage.measure import regionprops
+from common import load_segmentations, get_seg_path, get_all_tomograms, get_colormaps, STRUCTURE_NAMES
+
+
+def _create_pool_layer(seg, assignment_path):
+    assignments = pd.read_csv(assignment_path)
+    pool_names = pd.unique(assignments.pool).tolist()
+    pools = np.zeros_like(seg)
+
+    pool_colors = get_colormaps()["pools"]
+    colormap = {}
+    for pool_id, pool_name in enumerate(pool_names, 1):
+        pool_vesicle_ids = assignments[assignments.pool == pool_name].vesicle_id
+        pool_mask = np.isin(seg, pool_vesicle_ids)
+        pools[pool_mask] = pool_id
+        colormap[pool_id] = pool_colors[pool_name]
+
+    return pools, colormap
+
+
+def _update_assignments(vesicles, pool_correction, assignment_path):
+    old_assignments = pd.read_csv(assignment_path)
+    props = regionprops(vesicles, pool_correction)
+
+    new_assignments = old_assignments.copy()
+    val_to_pool = {1: "RA-V", 2: "MP-V", 3: "Docked-V", 4: None}
+    for prop in props:
+        correction_val = prop.max_intensity
+        if correction_val == 0:
+            continue
+        new_assignments[new_assignments.vesicle_id == prop.label] = val_to_pool[correction_val]
+
+    new_assignments.to_csv(assignment_path, index=False)
+
+
+# TODO: also enable correcting vesicle segmentation???
+def correct_vesicle_pools(mrc_path):
+    seg_path = get_seg_path(mrc_path)
+
+    output_folder = os.path.split(seg_path)[0]
+    assignment_path = os.path.join(output_folder, "vesicle_pools.csv")
+
+    data, _ = read_mrc(mrc_path)
+    segmentations = load_segmentations(seg_path)
+    vesicles = segmentations["vesicles"]
+
+    colormaps = get_colormaps()
+    pool_colors = colormaps["pools"]
+    correction_colors = {
+        1: pool_colors["RA-V"], 2: pool_colors["MP-V"], 3: pool_colors["Docked-V"], 4: "Gray", None: "Gray"
+    }
+
+    vesicle_pools, pool_colors = _create_pool_layer(vesicles, assignment_path)
+
+    pool_correction_path = os.path.join(output_folder, "correction", "pool_correction.tif")
+    if os.path.exists(pool_correction_path):
+        pool_correction = imageio.imread(pool_correction_path)
+    else:
+        pool_correction = np.zeros_like(vesicles)
+
+    v = napari.Viewer()
+    v.add_image(data)
+    v.add_labels(vesicle_pools, colormap=pool_colors)
+    v.add_labels(pool_correction, colormap=correction_colors)
+    v.add_labels(vesicles, visible=False)
+    for name in STRUCTURE_NAMES:
+        v.add_labels(segmentations[name], name=name, visible=False, colormap=colormaps[name])
+
+    @magicgui(call_button="Update Pools")
+    def update_pools(viewer: napari.Viewer):
+        pool_data = viewer.layers["vesicle_pools"].data
+        vesicles = viewer.layers["vesicles"].data
+        pool_correction = viewer.layers["pool_correction"].data
+        _update_assignments(vesicles, pool_correction, assignment_path)
+        imageio.imwrite(pool_correction_path, pool_correction, compression="zlib")
+        pool_data, pool_colors = _create_pool_layer(vesicles, assignment_path)
+        viewer.layers["vesicle_pools"].data = pool_data
+
+    v.window.add_dock_widget(update_pools)
+
+    napari.run()
+
+
+def main():
+    tomograms = get_all_tomograms()
+    for tomo in tomograms:
+        correct_vesicle_pools(tomo)
+
+
+if __name__ == "__main__":
+    main()