Merge branch 'main' of https://github.com/computational-cell-analytics/synaptic-reconstruction into sm-dev

Author: SarahMuth

scripts/baselines/cryo_ves_net/README.md

@@ -0,0 +1,5 @@
+# CryoVesNet
+
+Scripts to run CryoVesNet on our data. See https://github.com/Zuber-group/CryoVesNet for details.
+
+The code is currently not working due to this issue: https://github.com/Zuber-group/CryoVesNet/issues/6

scripts/baselines/cryo_ves_net/common.py

@@ -0,0 +1,123 @@
+import os
+import tempfile
+from glob import glob
+from pathlib import Path
+from shutil import copyfile
+
+import h5py
+import mrcfile
+import numpy as np
+
+import cryovesnet
+
+
+# additional parameters?
+def _segment_vesicles(directory):
+    pl = cryovesnet.Pipeline(directory, pattern="*.mrc")
+    pl.setup_cryovesnet_dir(make_masks=False)
+
+    pl.run_deep()
+    pl.rescale()
+    pl.label_vesicles(within_segmentation_region=False)
+    pl.label_vesicles_adaptive(separating=True)
+    pl.make_spheres()
+    pl.repair_spheres()
+
+
+def _prepare_input(path, output_folder, input_key, resolution, rel_folder=None):
+    fname = Path(path).stem
+    if rel_folder is None:
+        sub_folder = os.path.join(output_folder, fname)
+    else:
+        sub_folder = os.path.join(output_folder, rel_folder, fname)
+
+    os.makedirs(sub_folder, exist_ok=True)
+    out_path = os.path.join(sub_folder, f"{fname}.mrc")
+
+    if path.endswith(".h5"):
+        assert resolution is not None
+        with h5py.File(path, "r") as f:
+            vol = f[input_key][:]
+        mrcfile.new(out_path, data=vol)
+
+    # Copy the mrc file.
+    elif path.endswith(".mrc"):
+        copyfile(path, out_path)
+
+    # Update the resolution if it was given.
+    if resolution is not None:
+        with mrcfile.open(out_path, mode="r+") as f:
+            f.voxel_size = resolution
+            f.update_header_from_data()
+
+    return out_path, sub_folder
+
+
+def _process_output(tmp, tmp_file, output_folder, output_key, rel_folder=None, mask_file=None, mask_key=None):
+    fname = Path(tmp_file).stem
+    seg_path = os.path.join(tmp, "cryovesnet", f"{fname}_convex_labels.mrc")
+    with mrcfile.open(seg_path, "r") as f:
+        seg = f.data[:]
+
+    if mask_file is not None:
+        with h5py.File(mask_file, "r") as f:
+            mask = f[mask_key][:].astype("bool")
+        # We need to make this copy, otherwise seg is assignment only.
+        seg = np.asarray(seg).copy()
+        seg[~mask] = 0
+
+    if rel_folder is None:
+        this_output_folder = output_folder
+    else:
+        this_output_folder = os.path.join(output_folder, rel_folder)
+        os.makedirs(this_output_folder, exist_ok=True)
+
+    out_path = os.path.join(this_output_folder, f"{fname}.h5")
+    with h5py.File(out_path, "a") as f:
+        f.create_dataset(output_key, data=seg, compression="gzip")
+
+
+def apply_cryo_vesnet(
+    input_folder, output_folder, pattern, input_key,
+    resolution=None, output_key="prediction/vesicles/cryovesnet",
+    mask_folder=None, mask_key=None, nested=False,
+):
+    os.makedirs(output_folder, exist_ok=True)
+    if nested:
+        files = sorted(glob(os.path.join(input_folder, "**", pattern), recursive=True))
+    else:
+        files = sorted(glob(os.path.join(input_folder, pattern)))
+
+    if mask_folder is None:
+        mask_files = None
+    else:
+        assert mask_key is not None
+        if nested:
+            mask_files = sorted(glob(os.path.join(mask_folder, "**", pattern), recursive=True))
+        else:
+            mask_files = sorted(glob(os.path.join(mask_folder, pattern)))
+        assert len(mask_files) == len(files)
+
+    with tempfile.TemporaryDirectory() as tmp:
+
+        for i, file in enumerate(files):
+
+            # Get the resolution info for this file.
+            if resolution is None:
+                res = None
+            else:
+                fname = Path(file).stem
+                res = resolution[fname] if isinstance(resolution, dict) else resolution
+
+            # Prepare the input files by copying them over or resaving them (if h5).
+            rel_folder = os.path.split(os.path.relpath(file, input_folder))[0] if nested else None
+            tmp_file, sub_folder = _prepare_input(file, tmp, input_key, res, rel_folder=rel_folder)
+
+            # Segment the vesicles in the file.
+            _segment_vesicles(sub_folder)
+
+            # Write the output file.
+            mask_file = None if mask_files is None else mask_files[i]
+            _process_output(
+                sub_folder, tmp_file, output_folder, output_key, rel_folder, mask_file=mask_file, mask_key=mask_key
+            )

scripts/baselines/cryo_ves_net/segment_cooper.py

@@ -0,0 +1,14 @@
+from common import apply_cryo_vesnet
+
+
+def main():
+    input_folder = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/vesicles_processed_v2/testsets"  # noqa
+    output_folder = "./cryo-vesnet-test2"
+
+    # TODO determine the correct resolution (in angstrom) for each dataset
+    resolution = (10, 10, 10)
+    apply_cryo_vesnet(input_folder, output_folder, pattern="*.h5", input_key="raw", resolution=resolution, nested=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/baselines/cryo_ves_net/segment_cryo.py

@@ -0,0 +1,23 @@
+from common import apply_cryo_vesnet
+
+
+def main():
+    input_folder = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/fernandez-busnadiego/vesicle_gt/v2"  # noqa
+    output_folder = "./cryo-vesnet-test"
+
+    # Resolution in Angstrom in XYZ
+    # The two tomograms have a different resolution.
+    resolution = {
+        "vesicles-33K-L1": (14.6, 14.6, 14.6),
+        "vesicles-64K-LAM12": (7.56, 7.56, 7.56),
+    }
+    apply_cryo_vesnet(
+        input_folder, output_folder,
+        pattern="*.h5", input_key="raw",
+        mask_folder=input_folder, mask_key="/labels/mask",
+        resolution=resolution
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/cooper/export_vesicles_to_imod.py

@@ -5,7 +5,9 @@
 
 
 def export_vesicles_to_imod(args):
-    export_function = partial(write_segmentation_to_imod_as_points, min_radius=args.min_radius, radius_factor=args.increase_radius)
+    export_function = partial(
+        write_segmentation_to_imod_as_points, min_radius=args.min_radius, radius_factor=args.increase_radius
+    )
     export_helper(args.input_path, args.segmentation_path, args.output_path, export_function, force=args.force)
 
 
@@ -32,7 +34,7 @@ def main():
         help="Whether to over-write already present export results."
     )
     parser.add_argument(
-        "--increase_radius", type=float, default=1.5,
+        "--increase_radius", type=float, default=1.3,
         help="The factor to increase the radius of the exported vesicles.",
     )
     args = parser.parse_args()

scripts/cooper/ground_truth/compartments/export_compartment_gt.py

@@ -0,0 +1,106 @@
+import os
+from glob import glob
+from pathlib import Path
+
+import imageio.v3 as imageio
+import h5py
+import numpy as np
+
+from scipy.ndimage import binary_erosion, binary_closing
+from skimage.measure import label, regionprops
+from skimage.morphology import remove_small_holes
+from skimage.segmentation import watershed
+from synaptic_reconstruction.ground_truth.shape_refinement import edge_filter
+from tqdm import tqdm
+
+
+def process_compartment_gt(im_path, ann_path, output_root, view=True, snap_to_bd=False):
+    output_path = os.path.join(output_root, os.path.basename(im_path))
+    if os.path.exists(output_path):
+        return
+
+    seg = imageio.imread(ann_path)
+
+    with h5py.File(im_path, "r") as f:
+        tomo = f["data"][:]
+
+    if snap_to_bd:
+        hmap = edge_filter(tomo, sigma=3.0, per_slice=True)
+    else:
+        hmap = None
+
+    seg_pp = label(seg)
+    props = regionprops(seg_pp)
+
+    # for dilation / eroision
+    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
+
+    # Apply the post-processing for each segment.
+    min_size = 500
+    for prop in props:
+        # 1. size filter
+        if prop.area < min_size:
+            seg_pp[seg_pp == prop.label] = 0
+            continue
+
+        # 2. get the box and mask for the current object
+        bb = tuple(slice(start, stop) for start, stop in zip(prop.bbox[:3], prop.bbox[3:]))
+        mask = seg_pp[bb] == prop.label
+
+        # 3. filling smal holes and closing closing
+        mask = remove_small_holes(mask, area_threshold=500)
+        mask = np.logical_or(binary_closing(mask, iterations=4), mask)
+        mask = np.logical_or(binary_closing(mask, iterations=8, structure=structure_3d), mask)
+
+        # 4. snap to boundary
+        if snap_to_bd:
+            seeds = binary_erosion(mask, structure=structure_3d, iterations=3).astype("uint8")
+            bg_seeds = binary_erosion(~mask, structure=structure_3d, iterations=3)
+            seeds[bg_seeds] = 2
+            mask = watershed(hmap[bb], markers=seeds) == 1
+
+        # 5. write back
+        seg_pp[bb][mask] = prop.label
+
+    if view:
+        import napari
+
+        v = napari.Viewer()
+        v.add_image(tomo)
+        if hmap is not None:
+            v.add_image(hmap)
+        v.add_labels(seg, visible=False)
+        v.add_labels(seg_pp)
+        napari.run()
+        return
+
+    # Cut some border pixels to avoid artifacts.
+    bb = np.s_[4:-4, 16:-16, 16:-16]
+    tomo = tomo[bb]
+    seg_pp = seg_pp[bb]
+
+    with h5py.File(output_path, "a") as f:
+        f.create_dataset("raw", data=tomo, compression="gzip")
+        f.create_dataset("labels/compartments", data=seg_pp, compression="gzip")
+
+
+def main():
+    for ds in ["05_stem750_sv_training", "06_hoi_wt_stem750_fm"]:
+        annotation_folder = f"output/{ds}/segmentations"
+        annotations = sorted(glob(os.path.join(annotation_folder, "*.tif")))
+
+        output_root = f"output/compartment_gt/v2/{ds}"
+        os.makedirs(output_root, exist_ok=True)
+
+        image_folder = f"output/{ds}/tomograms"
+        for ann_path in tqdm(annotations):
+            fname = Path(ann_path).stem
+            im_path = os.path.join(image_folder, f"{fname}.h5")
+            assert os.path.exists(im_path)
+            process_compartment_gt(im_path, ann_path, output_root, view=False)
+
+
+if __name__ == "__main__":
+    main()

scripts/cooper/ground_truth/compartments/postprocess_annotations.py

@@ -0,0 +1,70 @@
+import os
+from glob import glob
+
+import h5py
+import imageio.v3 as imageio
+import napari
+import numpy as np
+
+from skimage.measure import label
+# from skimage.morphology import remove_small_holes
+from tqdm import tqdm
+
+
+def process_labels(labels):
+    labels = label(labels)
+
+    min_size = 75
+    ids, sizes = np.unique(labels, return_counts=True)
+    filter_ids = ids[sizes < min_size]
+    labels[np.isin(labels, filter_ids)] = 0
+
+    # labels = remove_small_holes(labels, area_threshold=min_size)
+    return labels
+
+
+def postprocess_annotation(im_path, ann_path, output_folder, view=False):
+    fname = os.path.basename(im_path)
+
+    out_path = os.path.join(output_folder, fname.replace(".tif", ".h5"))
+    if os.path.exists(out_path):
+        return
+
+    labels = imageio.imread(ann_path)
+
+    # Skip empty labels.
+    if labels.max() == 0:
+        print("Skipping", im_path)
+        return
+
+    image = imageio.imread(im_path)
+    labels = process_labels(labels)
+
+    if view:
+        v = napari.Viewer()
+        v.add_image(image)
+        v.add_labels(labels)
+        napari.run()
+        return
+
+    with h5py.File(out_path, "a") as f:
+        f.create_dataset("data", data=image, compression="gzip")
+        f.create_dataset("labels/compartments", data=labels, compression="gzip")
+
+
+def postprocess_annotations(view):
+    images = sorted(glob("output/images/*.tif"))
+    annotations = sorted(glob("output/annotations/*.tif"))
+
+    output_folder = "output/postprocessed_annotations"
+    os.makedirs(output_folder, exist_ok=True)
+    for im, ann in tqdm(zip(images, annotations), total=len(images)):
+        postprocess_annotation(im, ann, output_folder, view=view)
+
+
+def main():
+    postprocess_annotations(view=False)
+
+
+if __name__ == "__main__":
+    main()