Merge pull request #40 from computational-cell-analytics/compartment-training

Compartment segmentation and CryoVesNet

Author: constantinpape

scripts/baselines/cryo_ves_net/common.py

@@ -2,9 +2,12 @@
 import tempfile
 from glob import glob
 from pathlib import Path
+from shutil import copyfile
 
 import h5py
 import mrcfile
+import numpy as np
+
 import cryovesnet
 
 
@@ -14,51 +17,107 @@ def _segment_vesicles(directory):
     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):
-    out_path = os.path.join(output_folder, f"{Path(path).stem}.mrc")
+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)
-        with mrcfile.open(out_path, mode="r+") as f:
-            f.header.cella.x = resolution[0]
-            f.header.cella.y = resolution[1]
-            f.header.cella.z = resolution[2]
 
-    # TODO just copy the file
+    # Copy the mrc file.
     elif path.endswith(".mrc"):
-        pass
+        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")
 
 
-# TODO support nested
 def apply_cryo_vesnet(
     input_folder, output_folder, pattern, input_key,
-    resolution=None, output_key="prediction/vesicles/cryovesnet"
+    resolution=None, output_key="prediction/vesicles/cryovesnet",
+    mask_folder=None, mask_key=None, nested=False,
 ):
-    files = sorted(glob(os.path.join(input_folder, pattern)))
+    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:
 
-        # Prepare the input files by copying them over or resaving them (if h5).
-        for file in files:
+        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_input(file, tmp, input_key, res)
 
-        # Segment the vesicles in all files.
-        _segment_vesicles(tmp)
-        breakpoint()
+            # 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)
 
-        # TODO
-        # Re-save the segmentations to the output 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

@@ -11,7 +11,12 @@ def main():
         "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", resolution=resolution)
+    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__":

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/preprocess.py

@@ -40,8 +40,9 @@ def preprocess_tomogram(dataset, tomogram):
         output_path=output_path,
         model_type="vit_b",
         key="data",
-        checkpoint_path="./checkpoints/compartment_model/best.pt",
+        checkpoint_path="./checkpoints/compartment_model_v2/best.pt",
         ndim=3,
+        precompute_amg_state=True,
     )
 
 
@@ -78,7 +79,7 @@ def preprocess_cryo_tomogram(fname):
         output_path=output_path,
         model_type="vit_b",
         key="data",
-        checkpoint_path="./checkpoints/compartment_model/best.pt",
+        checkpoint_path="./checkpoints/compartment_model_v2/best.pt",
         ndim=3,
     )
 
@@ -113,10 +114,10 @@ def preprocess_cryo():
 
 
 def main():
-    # preprocess_05()
-    # preprocess_06()
-    # preprocess_09()
     preprocess_cryo()
+    preprocess_05()
+    preprocess_06()
+    preprocess_09()
 
 
 if __name__ == "__main__":

scripts/cooper/ground_truth/compartments/run_annotation.py

@@ -1,11 +1,13 @@
 import os
+from glob import glob
+from pathlib import Path
 
 from elf.io import open_file
 from micro_sam.sam_annotator import annotator_3d, image_folder_annotator
 
 
 def run_volume_annotation(ds, name):
-    checkpoint_path = "./checkpoints/compartment_model/best.pt"
+    checkpoint_path = "./checkpoints/compartment_model_v2/best.pt"
 
     tomogram_path = f"./output/{ds}/tomograms/{name}.h5"
     embedding_path = f"./output/{ds}/embeddings/{name}.zarr"
@@ -25,16 +27,40 @@ def run_image_annotation():
     )
 
 
-def main():
-    run_image_annotation()
+def annotate_cryo():
+    ds = "cryo"
+    name = "vesicles-33K-L1"
+    run_volume_annotation(ds, name)
+
+
+def _series_annotation(ds):
+    # name = "upSTEM750_36859_J2_TS_SP_001_rec_2kb1dawbp_crop"
+    images = glob(f"./output/{ds}/tomograms/*.h5")
+    for image in images:
+        name = Path(image).stem
+        seg_path = f"./output/{ds}/segmentations/{name}.tif"
+        print("Run segmentation for:", ds, name)
+        if os.path.exists(seg_path):
+            print("Skipping", ds, name, "because it is already segmented.")
+            continue
+        run_volume_annotation(ds, name)
 
-    # ds = "09_stem750_66k"
-    # name = "36859_J1_66K_TS_PS_05_rec_2kb1dawbp_crop"
 
-    # ds = "cryo"
-    # name = "vesicles-64K-LAM12"
+def annotate_05():
+    ds = "05_stem750_sv_training"
+    _series_annotation(ds)
 
-    # run_annotation(ds, name)
+
+def annotate_06():
+    ds = "06_hoi_wt_stem750_fm"
+    _series_annotation(ds)
+
+
+def main():
+    # run_image_annotation()
+    # annotate_cryo()
+    # annotate_05()
+    annotate_06()
 
 
 if __name__ == "__main__":