From ce94ea8018dd226d74bc9d5b71b0ccd3f4510741 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Mon, 25 Nov 2024 22:49:23 +0100 Subject: [PATCH] Add CLI for IMOD export --- setup.py | 2 + synaptic_reconstruction/imod/to_imod.py | 43 ++++++++++---- synaptic_reconstruction/tools/cli.py | 76 +++++++++++++++++++++++++ 3 files changed, 110 insertions(+), 11 deletions(-) diff --git a/setup.py b/setup.py index 4b898b3..3722b0d 100644 --- a/setup.py +++ b/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", diff --git a/synaptic_reconstruction/imod/to_imod.py b/synaptic_reconstruction/imod/to_imod.py index c9d48e7..8c30be9 100644 --- a/synaptic_reconstruction/imod/to_imod.py +++ b/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,6 +36,7 @@ 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. @@ -32,15 +44,14 @@ def write_segmentation_to_imod( 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) diff --git a/synaptic_reconstruction/tools/cli.py b/synaptic_reconstruction/tools/cli.py index e6482a7..bcb3085 100644 --- a/synaptic_reconstruction/tools/cli.py +++ b/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():