Update tomo sample data and example analysis pipeline

examples/analysis_pipeline.py

@@ -1,8 +1,14 @@
 import napari
 import pandas as pd
+import numpy as np
+
+from scipy.ndimage import binary_closing
+from skimage.measure import regionprops
+from skimage.segmentation import find_boundaries
 
 from synapse_net.file_utils import read_mrc
 from synapse_net.sample_data import get_sample_data
+from synapse_net.distance_measurements import measure_segmentation_to_object_distances
 from synapse_net.inference import compute_scale_from_voxel_size, get_model, run_segmentation
 
 
@@ -30,28 +36,77 @@ def segment_structures(tomogram, voxel_size):
     return {"vesicles": vesicles, "active_zone": active_zone, "compartments": compartments}
 
 
+def n_vesicles(mask, ves):
+    return len(np.unique(ves[mask])) - 1
+
+
 def postprocess_segmentation(segmentations):
-    pass
+    # We find the compartment corresponding to the presynaptic terminal
+    # by selecting the compartment with most vesicles. We filter out all
+    # vesicles that do not overlap with this compartment.
 
+    vesicles, compartments = segmentations["vesicles"], segmentations["compartments"]
 
-def measure_distances(segmentations):
-    pass
+    # First, we find the compartment with most vesicles.
+    props = regionprops(compartments, intensity_image=vesicles, extra_properties=[n_vesicles])
+    compartment_ids = [prop.label for prop in props]
+    vesicle_counts = [prop.n_vesicles for prop in props]
+    compartments = (compartments == compartment_ids[np.argmax(vesicle_counts)]).astype("uint8")
 
+    # Filter all vesicles that are not in the compartment.
+    props = regionprops(vesicles, compartments)
+    filter_ids = [prop.label for prop in props if prop.max_intensity == 0]
+    vesicles[np.isin(vesicles, filter_ids)] = 0
 
-def assign_vesicle_pools(distances):
-    pass
+    segmentations["vesicles"], segmentations["compartments"] = vesicles, compartments
+
+    # We also apply closing to the active zone segmentation to avoid gaps and then
+    # intersect it with the boundary of the presynaptic compartment.
+    active_zone = segmentations["active_zone"]
+    active_zone = binary_closing(active_zone, iterations=4)
+    boundary = find_boundaries(compartments)
+    active_zone = np.logical_and(active_zone, boundary).astype("uint8")
+    segmentations["active_zone"] = active_zone
+
+    return segmentations
 
 
-def visualize_results(tomogram, segmentations, vesicle_pools):
-    # TODO vesicle pool visualization
+def measure_distances(segmentations, voxel_size):
+    vesicles, active_zone = segmentations["vesicles"], segmentations["active_zone"]
+    voxel_size = tuple(voxel_size[ax] for ax in "zyx")
+    distances, _, _, vesicle_ids = measure_segmentation_to_object_distances(
+        vesicles, active_zone, resolution=voxel_size
+    )
+    return pd.DataFrame({"vesicle_id": vesicle_ids, "distance": distances})
+
+
+def assign_vesicle_pools(vesicle_attributes):
+    docked_vesicle_distance = 2  # nm
+    vesicle_attributes["pool"] = vesicle_attributes["distance"].apply(
+        lambda x: "docked" if x < docked_vesicle_distance else "non-attached"
+    )
+    return vesicle_attributes
+
+
+def visualize_results(tomogram, segmentations, vesicle_attributes):
+
+    # Create a segmentation to visualize the vesicle pools.
+    docked_ids = vesicle_attributes[vesicle_attributes.pool == "docked"].vesicle_id
+    non_attached_ids = vesicle_attributes[vesicle_attributes.pool == "non-attached"].vesicle_id
+    vesicles = segmentations["vesicles"]
+    vesicle_pools = np.isin(vesicles, docked_ids).astype("uint8")
+    vesicle_pools[np.isin(vesicles, non_attached_ids)] = 2
+
     viewer = napari.Viewer()
     viewer.add_image(tomogram)
     for name, segmentation in segmentations.items():
         viewer.add_labels(segmentation, name=name)
+    viewer.add_labels(vesicle_pools)
     napari.run()
 
 
-def save_analysis(segmentations, distances, vesicle_pools, save_path):
+# TODO compute the vesicle radii and other features and then save the attributes.
+def save_analysis(segmentations, vesicle_attributes, save_path):
     pass
 
 
@@ -64,28 +119,33 @@ def main():
     tomogram, voxel_size = read_mrc(mrc_path)
 
     # Segment synaptic vesicles, the active zone, and the synaptic compartment.
-    segmentations = segment_structures(tomogram, voxel_size)
+    # segmentations = segment_structures(tomogram, voxel_size)
+
+    # Load saved segmentations for development.
     import h5py
+    segmentations = {}
     with h5py.File("seg.h5", "r") as f:
-        for name, seg in segmentations.items():
-            f.create_dataset(name, data=seg, compression="gzip")
+        for name, ds in f.items():
+            # f.create_dataset(name, data=seg, compression="gzip")
+            seg = ds[:]
+            segmentations[name] = seg
 
     # Post-process the segmentations, to find the presynaptic terminal,
     # filter out vesicles not in the terminal, and to 'snape' the AZ to the presynaptic boundary.
     segmentations = postprocess_segmentation(segmentations)
 
     # Measure the distances between the AZ and vesicles.
-    distances = measure_distances(segmentations)
+    vesicle_attributes = measure_distances(segmentations, voxel_size)
 
     # Assign the vesicle pools, 'docked' and 'non-attached' vesicles, based on the distances.
-    vesicle_pools = assign_vesicle_pools(distances)
+    vesicle_attributes = assign_vesicle_pools(vesicle_attributes)
 
     # Visualize the results.
-    visualize_results(tomogram, segmentations, vesicle_pools)
+    visualize_results(tomogram, segmentations, vesicle_attributes)
 
     # Compute the vesicle radii and combine and save all measurements.
     save_path = "analysis_results.xlsx"
-    save_analysis(segmentations, distances, vesicle_pools, save_path)
+    save_analysis(segmentations, vesicle_attributes, save_path)
 
 
 if __name__ == "__main__":

synapse_net/sample_data.py

@@ -15,11 +15,11 @@ def get_sample_data(name: str) -> str:
     """
     registry = {
         "tem_2d.mrc": "3c6f9ff6d7673d9bf2fd46c09750c3c7dbb8fa1aa59dcdb3363b65cc774dcf28",
-        "tem_tomo.mrc": "eb790f83efb4c967c96961239ae52578d95da902fc32307629f76a26c3dc61fa",
+        "tem_tomo.mrc": "fe862ce7c22000d4440e3aa717ca9920b42260f691e5b2ab64cd61c928693c99",
     }
     urls = {
         "tem_2d.mrc": "https://owncloud.gwdg.de/index.php/s/5sAQ0U4puAspcHg/download",
-        "tem_tomo.mrc": "https://owncloud.gwdg.de/index.php/s/TmLjDCXi42E49Ef/download",
+        "tem_tomo.mrc": "https://owncloud.gwdg.de/index.php/s/FJDhDfbT4UxhtOn/download",
     }
     key = f"{name}.mrc"