Add CLI for IMOD export

Author: constantinpape

setup.py

@@ -14,6 +14,8 @@
     entry_points={
         "console_scripts": [
             "synapse_net.run_segmentation = synaptic_reconstruction.tools.cli:segmentation_cli",
+            "synapse_net.export_to_imod_points = synaptic_reconstruction.tools.cli:imod_point_cli",
+            "synapse_net.export_to_imod_objects = synaptic_reconstruction.tools.cli:imod_object_cli",
         ],
         "napari.manifest": [
             "synaptic_reconstruction = synaptic_reconstruction:napari.yaml",

synaptic_reconstruction/imod/to_imod.py

@@ -8,6 +8,7 @@
 from subprocess import run
 from typing import Optional, Tuple, Union
 
+import h5py
 import imageio.v3 as imageio
 import mrcfile
 import numpy as np
@@ -16,6 +17,16 @@
 from tqdm import tqdm
 
 
+def _load_segmentation(segmentation_path, segmentation_key):
+    assert os.path.exists(segmentation_path), segmentation_path
+    if segmentation_key is None:
+        seg = imageio.imread(segmentation_path)
+    else:
+        with h5py.File(segmentation_path, "r") as f:
+            seg = f[segmentation_key][:]
+    return seg
+
+
 # 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.
@@ -25,22 +36,22 @@ def write_segmentation_to_imod(
     mrc_path: str,
     segmentation: Union[str, np.ndarray],
     output_path: str,
+    segmentation_key: Optional[str] = None,
 ) -> None:
     """Write a segmentation to a mod file as closed contour objects.
 
     Args:
         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.
+        segmentation_key: The key to the segmentation data in case the segmentation is stored in hdf5 files.
     """
     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)
+    # Load the segmentation case a filepath was passed.
+    segmentation = _load_segmentation(segmentation, segmentation_key)
 
     # Binarize the segmentation and flip its axes to match the IMOD axis convention.
     segmentation = (segmentation > 0).astype("uint8")
@@ -187,6 +198,7 @@ def write_segmentation_to_imod_as_points(
     min_radius: Union[int, float],
     radius_factor: float = 1.0,
     estimate_radius_2d: bool = True,
+    segmentation_key: Optional[str] = None,
 ) -> None:
     """Write segmentation results to .mod file with imod point annotations.
 
@@ -201,6 +213,7 @@ def write_segmentation_to_imod_as_points(
         estimate_radius_2d: If true the distance to boundary for determining the centroid and computing
             the radius will be computed only in 2d rather than in 3d. This can lead to better results
             in case of deformation across the depth axis.
+        segmentation_key: The key to the segmentation data in case the segmentation is stored in hdf5 files.
     """
 
     # Read the resolution information from the mrcfile.
@@ -212,7 +225,7 @@ def write_segmentation_to_imod_as_points(
 
     # Extract the center coordinates and radii from the segmentation.
     if isinstance(segmentation, str):
-        segmentation = imageio.imread(segmentation)
+        segmentation = _load_segmentation(segmentation, segmentation_key)
     coordinates, radii = convert_segmentation_to_spheres(
         segmentation, resolution=resolution, radius_factor=radius_factor, estimate_radius_2d=estimate_radius_2d
     )
@@ -221,16 +234,22 @@ def write_segmentation_to_imod_as_points(
     write_points_to_imod(coordinates, radii, segmentation.shape, min_radius, output_path)
 
 
-# TODO we also need to support .rec files ...
-def _get_file_paths(input_path, ext=".mrc"):
+def _get_file_paths(input_path, ext=(".mrc", ".rec")):
     if not os.path.exists(input_path):
-        raise Exception(f"Input path not found {input_path}")
+        raise Exception(f"Input path not found {input_path}.")
+
+    if isinstance(ext, str):
+        ext = (ext,)
 
     if os.path.isfile(input_path):
         input_files = [input_path]
         input_root = None
     else:
-        input_files = sorted(glob(os.path.join(input_path, "**", f"*{ext}"), recursive=True))
+        input_files = []
+        for ex in ext:
+            input_files.extend(
+                sorted(glob(os.path.join(input_path, "**", f"*{ex}"), recursive=True))
+            )
         input_root = input_path
 
     return input_files, input_root
@@ -242,6 +261,7 @@ def export_helper(
     output_root: str,
     export_function: callable,
     force: bool = False,
+    segmentation_key: Optional[str] = None,
 ) -> None:
     """
     Helper function to run imod export for files in a directory.
@@ -258,9 +278,10 @@ def export_helper(
             the path to the segmentation in a .tif file and the output path as only arguments.
             If you want to pass additional arguments to this function the use 'funtools.partial'
         force: Whether to rerun segmentation for output files that are already present.
+        segmentation_key: The key to the segmentation data in case the segmentation is stored in hdf5 files.
     """
     input_files, input_root = _get_file_paths(input_path)
-    segmentation_files, _ = _get_file_paths(segmentation_path, ext=".tif")
+    segmentation_files, _ = _get_file_paths(segmentation_path, ext=".tif" if segmentation_key is None else ".h5")
     assert len(input_files) == len(segmentation_files)
 
     for input_path, seg_path in tqdm(zip(input_files, segmentation_files), total=len(input_files)):
@@ -279,4 +300,4 @@ def export_helper(
             continue
 
         os.makedirs(os.path.split(output_path)[0], exist_ok=True)
-        export_function(input_path, seg_path, output_path)
+        export_function(input_path, seg_path, output_path, segmentation_key=segmentation_key)

synaptic_reconstruction/tools/cli.py

@@ -2,9 +2,85 @@
 from functools import partial
 
 from .util import run_segmentation, get_model
+from ..imod.to_imod import export_helper, write_segmentation_to_imod_as_points, write_segmentation_to_imod
 from ..inference.util import inference_helper, parse_tiling
 
 
+def imod_point_cli():
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument(
+        "--input_path", "-i", required=True,
+        help="The filepath to the mrc file or the directory containing the tomogram data."
+    )
+    parser.add_argument(
+        "--segmentation_path", "-s", required=True,
+        help="The filepath to the tif file or the directory containing the segmentations."
+    )
+    parser.add_argument(
+        "--output_path", "-o", required=True,
+        help="The filepath to directory where the segmentations will be saved."
+    )
+    parser.add_argument(
+        "--segmentation_key", "-k", help=""
+    )
+    parser.add_argument(
+        "--min_radius", type=float, default=10.0, help=""
+    )
+    parser.add_argument(
+        "--radius_factor", type=float, default=1.0, help="",
+    )
+    parser.add_argument(
+        "--force", action="store_true", help="",
+    )
+    args = parser.parse_args()
+
+    export_function = partial(
+        write_segmentation_to_imod_as_points,
+        min_radius=args.min_radius,
+        radius_factor=args.radius_factor,
+    )
+
+    export_helper(
+        input_path=args.input_path,
+        segmentation_path=args.segmentation_path,
+        output_root=args.output_path,
+        export_function=export_function,
+        force=args.force,
+        segmentation_key=args.segmentation_key,
+    )
+
+
+def imod_object_cli():
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument(
+        "--input_path", "-i", required=True,
+        help="The filepath to the mrc file or the directory containing the tomogram data."
+    )
+    parser.add_argument(
+        "--segmentation_path", "-s", required=True,
+        help="The filepath to the tif file or the directory containing the segmentations."
+    )
+    parser.add_argument(
+        "--output_path", "-o", required=True,
+        help="The filepath to directory where the segmentations will be saved."
+    )
+    parser.add_argument(
+        "--segmentation_key", "-k", help=""
+    )
+    parser.add_argument(
+        "--force", action="store_true", help="",
+    )
+    args = parser.parse_args()
+    export_helper(
+        input_path=args.input_path,
+        segmentation_path=args.segmentation_path,
+        output_root=args.output_path,
+        export_function=write_segmentation_to_imod,
+        force=args.force,
+        segmentation_key=args.segmentation_key,
+    )
+
+
 # TODO: handle kwargs
 # TODO: add custom model path
 def segmentation_cli():