Fix issues with the segmentation export to IMOD

Author: constantinpape

scripts/cooper/export_mask_to_imod.py

@@ -4,19 +4,11 @@
 
 
 def export_mask_to_imod(args):
-    # Test script
-    # write_segmentation_to_imod(
-    #     "synapse-examples/36859_J1_66K_TS_CA3_PS_26_rec_2Kb1dawbp_crop.mrc",
-    #     "synapse-examples/36859_J1_66K_TS_CA3_PS_26_rec_2Kb1dawbp_crop_mitos.tif",
-    #     "synapse-examples/mito.mod"
-    # )
     write_segmentation_to_imod(args.input_path, args.segmentation_path, args.output_path)
 
 
 def main():
     parser = argparse.ArgumentParser()
-
-    args = parser.parse_args()
     parser.add_argument(
         "-i", "--input_path", required=True,
         help="The filepath to the mrc file containing the data."

scripts/inner_ear/analysis/.gitignore

@@ -1,2 +1,3 @@
 panels/
+auto_seg_export/
 *.zip

scripts/inner_ear/analysis/export_seg_to_imod.py

@@ -0,0 +1,128 @@
+import os
+from shutil import copyfile
+from subprocess import run
+
+import imageio.v3 as imageio
+import mrcfile
+import napari
+import numpy as np
+import pandas as pd
+from elf.io import open_file
+from skimage.transform import resize
+from synaptic_reconstruction.imod.to_imod import write_segmentation_to_imod, write_segmentation_to_imod_as_points
+
+out_folder = "./auto_seg_export"
+os.makedirs(out_folder, exist_ok=True)
+
+
+def _resize(seg, tomo_path):
+    with open_file(tomo_path, "r") as f:
+        shape = f["data"].shape
+
+    if shape != seg.shape:
+        seg = resize(seg, shape, order=0, anti_aliasing=False, preserve_range=True).astype(seg.dtype)
+    assert seg.shape == shape
+    return seg
+
+
+def check_imod(tomo_path, mod_path):
+    run(["imod", tomo_path, mod_path])
+
+
+def export_pool(pool_name, pool_seg, tomo_path):
+    seg_path = f"./auto_seg_export/{pool_name}.tif"
+    pool_seg = _resize(pool_seg, tomo_path)
+    imageio.imwrite(seg_path, pool_seg, compression="zlib")
+
+    output_path = f"./auto_seg_export/{pool_name}.mod"
+    write_segmentation_to_imod_as_points(tomo_path, seg_path, output_path, min_radius=5)
+
+    check_imod(tomo_path, output_path)
+
+
+def export_vesicles(folder, tomo_path):
+    vesicle_pool_path = os.path.join(folder, "Korrektur", "vesicle_pools.tif")
+    # pool_correction_path = os.path.join(folder, "Korrektur", "pool_correction.tif")
+    # pool_correction = imageio.imread(pool_correction_path)
+
+    assignment_path = os.path.join(folder, "Korrektur", "measurements.xlsx")
+    assignments = pd.read_excel(assignment_path)
+
+    vesicles = imageio.imread(vesicle_pool_path)
+
+    pools = {}
+    for pool_name in pd.unique(assignments.pool):
+        pool_ids = assignments[assignments.pool == pool_name].id.values
+        pool_seg = vesicles.copy()
+        pool_seg[~np.isin(vesicles, pool_ids)] = 0
+        pools[pool_name] = pool_seg
+
+    view = False
+    if view:
+        v = napari.Viewer()
+        v.add_labels(vesicles, visible=False)
+        for pool_name, pool_seg in pools.items():
+            v.add_labels(pool_seg, name=pool_name)
+        napari.run()
+    else:
+        for pool_name, pool_seg in pools.items():
+            export_pool(pool_name, pool_seg, tomo_path)
+
+
+def export_structure(folder, tomo, name, view=False):
+    path = os.path.join(folder, "Korrektur", f"{name}.tif")
+    seg = imageio.imread(path)
+    seg = _resize(seg, tomo)
+
+    if view:
+        with open_file(tomo, "r") as f:
+            raw = f["data"][:]
+
+        v = napari.Viewer()
+        v.add_image(raw)
+        v.add_labels(seg)
+        napari.run()
+
+        return
+
+    seg_path = f"./auto_seg_export/{name}.tif"
+    imageio.imwrite(seg_path, seg, compression="zlib")
+    output_path = f"./auto_seg_export/{name}.mod"
+    write_segmentation_to_imod(tomo, seg_path, output_path)
+    check_imod(tomo, output_path)
+
+
+def remove_scale(tomo):
+    new_path = "./auto_seg_export/Emb71M1aGridA1sec1mod7.rec.rec"
+    if os.path.exists(new_path):
+        return new_path
+
+    copyfile(tomo, new_path)
+
+    with mrcfile.open(new_path, "r+") as f:
+        # Set the origin to (0, 0, 0)
+        f.header.nxstart = 0
+        f.header.nystart = 0
+        f.header.nzstart = 0
+        f.header.origin = (0.0, 0.0, 0.0)
+
+        # Save changes
+        f.flush()
+
+    return new_path
+
+
+def main():
+    folder = "/home/pape/Work/data/moser/em-synapses/Electron-Microscopy-Susi/Analyse/WT strong stim/Mouse 1/modiolar/1"
+    tomo = os.path.join(folder, "Emb71M1aGridA1sec1mod7.rec.rec")
+
+    tomo = remove_scale(tomo)
+
+    # export_vesicles(folder, tomo)
+    # export_structure(folder, tomo, "ribbon", view=False)
+    # export_structure(folder, tomo, "membrane", view=False)
+    export_structure(folder, tomo, "PD", view=False)
+
+
+if __name__ == "__main__":
+    main()

synaptic_reconstruction/imod/to_imod.py

@@ -16,51 +16,60 @@
 from tqdm import tqdm
 
 
-# FIXME how to bring the data to the IMOD axis convention?
-def _to_imod_order(data):
-    # data = np.swapaxes(data, 0, -1)
-    # data = np.fliplr(data)
-    # data = np.swapaxes(data, 0, -1)
-    return data
-
-
+# TODO: this has still some issues with some tomograms that has an offset info.
+# For now, this occurs for the inner ear data tomograms; it works for Fidi's STEM tomograms.
+# Ben's theory is that this might be due to data form JEOL vs. ThermoFischer microscopes.
+# To test this I can check how it works for data from Maus et al. / Imig et al., which were taken on a JEOL.
+# Can also check out the mrc documentation here: https://www.ccpem.ac.uk/mrc_format/mrc2014.php
 def write_segmentation_to_imod(
     mrc_path: str,
-    segmentation_path: str,
+    segmentation: Union[str, np.ndarray],
     output_path: str,
 ) -> None:
-    """Write a segmentation to a mod file as contours.
+    """Write a segmentation to a mod file as closed contour objects.
 
     Args:
-        mrc_path: a
-        segmentation_path: a
-        output_path: a
+        mrc_path: The filepath to the mrc file from which the segmentation was derived.
+        segmentation: The segmentation (either as numpy array or filepath to a .tif file).
+        output_path: The output path where the mod file will be saved.
     """
     cmd = "imodauto"
     cmd_path = shutil.which(cmd)
     assert cmd_path is not None, f"Could not find the {cmd} imod command."
 
+    # Load the segmentation from a tif file in case a filepath was passed.
+    if isinstance(segmentation, str):
+        assert os.path.exists(segmentation)
+        segmentation = imageio.imread(segmentation)
+
+    # Binarize the segmentation and flip its axes to match the IMOD axis convention.
+    segmentation = (segmentation > 0).astype("uint8")
+    segmentation = np.flip(segmentation, axis=1)
+
+    # Read the voxel size and origin information from the mrc file.
     assert os.path.exists(mrc_path)
-    with mrcfile.open(mrc_path, mode="r+") as f:
+    with mrcfile.open(mrc_path, mode="r") as f:
         voxel_size = f.voxel_size
+        nx, ny, nz = f.header.nxstart, f.header.nystart, f.header.nzstart
+        origin = f.header.origin
 
+    # Write the input for imodauto to a temporary mrc file.
     with tempfile.NamedTemporaryFile(suffix=".mrc") as f:
         tmp_path = f.name
 
-        seg = (imageio.imread(segmentation_path) > 0).astype("uint8")
-        seg_ = _to_imod_order(seg)
-
-        # import napari
-        # v = napari.Viewer()
-        # v.add_image(seg)
-        # v.add_labels(seg_)
-        # napari.run()
-
-        mrcfile.new(tmp_path, data=seg_, overwrite=True)
+        mrcfile.new(tmp_path, data=segmentation, overwrite=True)
+        # Write the voxel_size and origin infomration.
         with mrcfile.open(tmp_path, mode="r+") as f:
             f.voxel_size = voxel_size
+
+            f.header.nxstart = nx
+            f.header.nystart = ny
+            f.header.nzstart = nz
+            f.header.origin = (0.0, 0.0, 0.0) * 3 if origin is None else origin
+
             f.update_header_from_data()
 
+        # Run the command.
         cmd_list = [cmd, "-E", "1", "-u", tmp_path, output_path]
         run(cmd_list)