synapse_net.tools.cli

  1import argparse
  2import os
  3from functools import partial
  4
  5import torch
  6import torch_em
  7from tqdm import tqdm
  8
  9from ..cristae_analysis import compute_mito_crista_statistics
 10from ..file_utils import read_voxel_size
 11from ..imod.to_imod import (
 12    _get_file_paths, _load_segmentation, export_helper,
 13    write_segmentation_to_imod_as_points, write_segmentation_to_imod,
 14)
 15from ..inference.inference import _get_model_registry, get_model, get_model_training_resolution, run_segmentation
 16from ..inference.scalable_segmentation import scalable_segmentation
 17from ..inference.util import inference_helper, parse_tiling
 18from .pool_visualization import _visualize_vesicle_pools
 19
 20
 21def imod_point_cli():
 22    parser = argparse.ArgumentParser(
 23        description="Convert a vesicle segmentation to an IMOD point model, "
 24        "corresponding to a sphere for each vesicle in the segmentation."
 25    )
 26    parser.add_argument(
 27        "--input_path", "-i", required=True,
 28        help="The filepath to the mrc file or the directory containing the tomogram data."
 29    )
 30    parser.add_argument(
 31        "--segmentation_path", "-s", required=True,
 32        help="The filepath to the file or the directory containing the segmentations."
 33    )
 34    parser.add_argument(
 35        "--output_path", "-o", required=True,
 36        help="The filepath to directory where the segmentations will be saved."
 37    )
 38    parser.add_argument(
 39        "--segmentation_key", "-k",
 40        help="The key in the segmentation files. If not given we assume that the segmentations are stored as tif."
 41        "If given, we assume they are stored as hdf5 files, and use the key to load the internal dataset."
 42    )
 43    parser.add_argument(
 44        "--min_radius", type=float, default=10.0,
 45        help="The minimum vesicle radius in nm. Objects that are smaller than this radius will be exclded from the export."  # noqa
 46    )
 47    parser.add_argument(
 48        "--radius_factor", type=float, default=1.0,
 49        help="A factor for scaling the sphere radius for the export. "
 50        "This can be used to fit the size of segmented vesicles to the best matching spheres.",
 51    )
 52    parser.add_argument(
 53        "--force", action="store_true",
 54        help="Whether to over-write already present export results."
 55    )
 56    args = parser.parse_args()
 57
 58    export_function = partial(
 59        write_segmentation_to_imod_as_points,
 60        min_radius=args.min_radius,
 61        radius_factor=args.radius_factor,
 62    )
 63
 64    export_helper(
 65        input_path=args.input_path,
 66        segmentation_path=args.segmentation_path,
 67        output_root=args.output_path,
 68        export_function=export_function,
 69        force=args.force,
 70        segmentation_key=args.segmentation_key,
 71    )
 72
 73
 74def imod_object_cli():
 75    parser = argparse.ArgumentParser(
 76        description="Convert segmented objects to close contour IMOD models."
 77    )
 78    parser.add_argument(
 79        "--input_path", "-i", required=True,
 80        help="The filepath to the mrc file or the directory containing the tomogram data."
 81    )
 82    parser.add_argument(
 83        "--segmentation_path", "-s", required=True,
 84        help="The filepath to the file or the directory containing the segmentations."
 85    )
 86    parser.add_argument(
 87        "--output_path", "-o", required=True,
 88        help="The filepath to directory where the segmentations will be saved."
 89    )
 90    parser.add_argument(
 91        "--segmentation_key", "-k",
 92        help="The key in the segmentation files. If not given we assume that the segmentations are stored as tif."
 93        "If given, we assume they are stored as hdf5 files, and use the key to load the internal dataset."
 94    )
 95    parser.add_argument(
 96        "--force", action="store_true",
 97        help="Whether to over-write already present export results."
 98    )
 99    args = parser.parse_args()
100    export_helper(
101        input_path=args.input_path,
102        segmentation_path=args.segmentation_path,
103        output_root=args.output_path,
104        export_function=write_segmentation_to_imod,
105        force=args.force,
106        segmentation_key=args.segmentation_key,
107    )
108
109
110def pool_visualization_cli():
111    parser = argparse.ArgumentParser(description="Load tomogram data, vesicle pools and additional segmentations for viualization.")  # noqa
112    parser.add_argument(
113        "--input_path", "-i", required=True,
114        help="The filepath to the mrc file containing the tomogram data."
115    )
116    parser.add_argument(
117        "--vesicle_paths", "-v", required=True, nargs="+",
118        help="The filepath(s) to the tif file(s) containing the vesicle segmentation."
119    )
120    parser.add_argument(
121        "--table_paths", "-t", required=True, nargs="+",
122        help="The filepath(s) to the table(s) with the vesicle pool assignments."
123    )
124    parser.add_argument(
125        "-s", "--segmentation_paths", nargs="+", help="Filepaths for additional segmentations."
126    )
127    parser.add_argument(
128        "--split_pools", action="store_true", help="Whether to split the pools into individual layers.",
129    )
130    args = parser.parse_args()
131    _visualize_vesicle_pools(
132        args.input_path, args.vesicle_paths, args.table_paths, args.segmentation_paths, args.split_pools
133    )
134
135
136# TODO: handle kwargs
137def segmentation_cli():
138    parser = argparse.ArgumentParser(description="Run segmentation.")
139    parser.add_argument(
140        "--input_path", "-i", required=True,
141        help="The filepath to the mrc file or the directory containing the tomogram data."
142    )
143    parser.add_argument(
144        "--output_path", "-o", required=True,
145        help="The filepath to directory where the segmentations will be saved."
146    )
147    model_names = list(_get_model_registry().urls.keys())
148    model_names = ", ".join(model_names)
149    parser.add_argument(
150        "--model", "-m", required=True,
151        help=f"The model type. The following models are currently available: {model_names}"
152    )
153    parser.add_argument(
154        "--mask_path", help="The filepath to a tif file with a mask that will be used to restrict the segmentation."
155        "Can also be a directory with tifs if the filestructure matches input_path."
156    )
157    parser.add_argument("--input_key", "-k", required=False)
158    parser.add_argument(
159        "--force", action="store_true",
160        help="Whether to over-write already present segmentation results."
161    )
162    parser.add_argument(
163        "--tile_shape", type=int, nargs=3,
164        help="The tile shape for prediction, in ZYX order. Lower the tile shape if GPU memory is insufficient."
165    )
166    parser.add_argument(
167        "--halo", type=int, nargs=3,
168        help="The halo for prediction, in ZYX order. Increase the halo to minimize boundary artifacts."
169    )
170    parser.add_argument(
171        "--data_ext", default=".mrc", help="The extension of the tomogram data. By default .mrc."
172    )
173    parser.add_argument(
174        "--checkpoint", "-c", help="Path to a custom model, e.g. from domain adaptation.",
175    )
176    parser.add_argument(
177        "--segmentation_key", "-s",
178        help="If given, the outputs will be saved to an hdf5 file with this key. Otherwise they will be saved as tif.",
179    )
180    parser.add_argument(
181        "--scale", type=float,
182        help="The factor for rescaling the data before inference. "
183        "By default, the scaling factor will be derived from the voxel size of the input data. "
184        "If this parameter is given it will over-ride the default behavior. "
185    )
186    parser.add_argument(
187        "--verbose", "-v", action="store_true",
188        help="Whether to print verbose information about the segmentation progress."
189    )
190    parser.add_argument(
191        "--scalable", action="store_true", help="Use the scalable segmentation implementation. "
192        "Currently this only works for vesicles, mitochondria, or active zones."
193    )
194    parser.add_argument(
195        "--extra_input_path", default=None, help="Filepath to extra inputs, needed for cristae segmentation."
196    )
197    parser.add_argument(
198        "--extra_input_ext", default=".tif", help="File extension for the extra inputs, default is tif."
199    )
200    args = parser.parse_args()
201
202    if args.checkpoint is None:
203        model = get_model(args.model)
204    else:
205        checkpoint_path = args.checkpoint
206        if checkpoint_path.endswith("best.pt"):
207            checkpoint_path = os.path.split(checkpoint_path)[0]
208
209        if os.path.isdir(checkpoint_path):  # Load the model from a torch_em checkpoint.
210            model = torch_em.util.load_model(checkpoint=checkpoint_path)
211        else:
212            model = torch.load(checkpoint_path, weights_only=False)
213        assert model is not None, f"The model from {args.checkpoint} could not be loaded."
214
215    is_2d = "2d" in args.model
216    tiling = parse_tiling(args.tile_shape, args.halo, is_2d=is_2d)
217
218    # If the scale argument is not passed, then we get the average training resolution for the model.
219    # The inputs will then be scaled to match this resolution based on the voxel size from the mrc files.
220    if args.scale is None:
221        model_resolution = get_model_training_resolution(args.model)
222        model_resolution = tuple(model_resolution[ax] for ax in ("yx" if is_2d else "zyx"))
223        scale = None
224    # Otherwise, we set the model resolution to None and use the scaling factor provided by the user.
225    else:
226        model_resolution = None
227        scale = (2 if is_2d else 3) * (args.scale,)
228
229    if args.scalable:
230        if not args.model.startswith(("vesicle", "mito", "active")):
231            raise ValueError(
232                "The scalable segmentation implementation is currently only supported for "
233                f"vesicles, mitochondria, or active zones, not for {args.model}."
234            )
235        segmentation_function = partial(
236            scalable_segmentation, model=model, tiling=tiling, verbose=args.verbose
237        )
238        allocate_output = True
239
240    else:
241        segmentation_function = partial(
242            run_segmentation, model=model, model_type=args.model, verbose=args.verbose, tiling=tiling,
243        )
244        allocate_output = False
245
246    inference_helper(
247        args.input_path, args.output_path, segmentation_function,
248        mask_input_path=args.mask_path, force=args.force, data_ext=args.data_ext,
249        output_key=args.segmentation_key, model_resolution=model_resolution, scale=scale,
250        allocate_output=allocate_output, extra_input_path=args.extra_input_path,
251        extra_input_ext=args.extra_input_ext
252    )
253
254
255def cristae_analysis_helper(
256    crista_path, mito_path, output_root,
257    crista_key=None, mito_key=None,
258    voxel_size=None, tomogram_path=None,
259    membrane_thickness_nm=8.0, border_gap_nm=None,
260    method="skip", membrane_mode="slice_2d",
261    n_jobs=-1, force=False, verbose=False,
262):
263    """Batch-compute per-mitochondrion cristae statistics and save one CSV per input pair.
264
265    This is the headless equivalent of the napari cristae-analysis widget. It matches crista and
266    mitochondria segmentations by sorted order (a single file each, or two directories), computes
267    the statistics via :func:`synapse_net.cristae_analysis.compute_mito_crista_statistics`, and
268    writes the resulting table next to a mirrored input folder structure.
269
270    Args:
271        crista_path: Crista segmentation - a single file or a directory of them.
272        mito_path: Mitochondria instance segmentation - a single file or a directory of them.
273        output_root: Directory where the ``<stem>_cristae_analysis.csv`` tables are written. A single
274            input file writes directly into it; a directory input mirrors the nested folder structure.
275        crista_key: Internal dataset key for the crista segmentation. If None the crista files are
276            assumed to be tif, otherwise hdf5 with this key.
277        mito_key: Internal dataset key for the mitochondria segmentation, analogous to crista_key.
278        voxel_size: Voxel size in nm applied to every file. If None it is read per file from the
279            raw tomogram given via tomogram_path.
280        tomogram_path: Raw tomogram (mrc/rec) - a single file or a directory - used to read the
281            voxel size when voxel_size is None.
282        membrane_thickness_nm: Membrane shell thickness in nm.
283        border_gap_nm: Distance from the volume faces where the membrane is suppressed (nm).
284            Defaults to membrane_thickness_nm when None.
285        method: How the crista orientation anisotropy is computed ("skip", "fast" or "exact").
286        membrane_mode: How the membrane shell is built ("slice_2d" or "shell_3d").
287        n_jobs: Number of workers for the per-mitochondrion computation (-1 = all cores).
288        force: Whether to over-write already present result tables.
289        verbose: Whether to show a progress bar over the mitochondria of each file.
290    """
291    crista_files, crista_root = _get_file_paths(crista_path, ext=".h5" if crista_key else ".tif")
292    mito_files, _ = _get_file_paths(mito_path, ext=".h5" if mito_key else ".tif")
293    if len(crista_files) != len(mito_files):
294        raise ValueError(
295            f"The number of crista ({len(crista_files)}) and mitochondria ({len(mito_files)}) "
296            "segmentations does not match."
297        )
298
299    if voxel_size is not None:
300        voxel_sizes = [voxel_size] * len(crista_files)
301    elif tomogram_path is not None:
302        tomo_files, _ = _get_file_paths(tomogram_path, ext=(".mrc", ".rec"))
303        if len(tomo_files) != len(crista_files):
304            raise ValueError(
305                f"The number of tomograms ({len(tomo_files)}) does not match the number of "
306                f"crista segmentations ({len(crista_files)})."
307            )
308        voxel_sizes = [read_voxel_size(path) for path in tomo_files]
309    else:
310        raise ValueError("Provide either --voxel_size or --tomogram_path to determine the voxel size.")
311
312    for crista_file, mito_file, this_voxel_size in tqdm(
313        zip(crista_files, mito_files, voxel_sizes), total=len(crista_files), desc="Processing files"
314    ):
315        input_folder, input_name = os.path.split(crista_file)
316        fname = os.path.splitext(input_name)[0] + "_cristae_analysis.csv"
317        if crista_root is None:
318            output_path = os.path.join(output_root, fname)
319        else:
320            rel_folder = os.path.relpath(input_folder, crista_root)
321            output_path = os.path.join(output_root, rel_folder, fname)
322
323        if os.path.exists(output_path) and not force:
324            continue
325
326        crista = _load_segmentation(crista_file, crista_key)
327        mito = _load_segmentation(mito_file, mito_key)
328        stats_df = compute_mito_crista_statistics(
329            crista, mito, this_voxel_size,
330            membrane_thickness_nm=membrane_thickness_nm, border_gap_nm=border_gap_nm,
331            method=method, membrane_mode=membrane_mode, n_jobs=n_jobs, verbose=verbose,
332        )
333
334        os.makedirs(os.path.split(output_path)[0], exist_ok=True)
335        stats_df.to_csv(output_path, index=False)
336        print(f"Saved cristae analysis to {output_path}.")
337
338
339def cristae_analysis_cli():
340    parser = argparse.ArgumentParser(
341        description="Compute per-mitochondrion cristae statistics from a crista segmentation and a "
342        "mitochondria instance segmentation, and save the results as a CSV table. This is the "
343        "command-line equivalent of the napari cristae-analysis widget."
344    )
345    parser.add_argument(
346        "--crista_path", "-c", required=True,
347        help="The filepath to the crista segmentation, or a directory containing multiple of them."
348    )
349    parser.add_argument(
350        "--mito_path", "-m", required=True,
351        help="The filepath to the mitochondria instance segmentation, or a directory containing multiple of them."
352    )
353    parser.add_argument(
354        "--output_path", "-o", required=True,
355        help="The filepath to the directory where the result tables will be saved."
356    )
357    parser.add_argument(
358        "--crista_key",
359        help="The key in the crista segmentation file. If not given the crista segmentation is assumed to be tif. "
360        "If given, it is assumed to be an hdf5 file and the key is used to load the internal dataset."
361    )
362    parser.add_argument(
363        "--mito_key",
364        help="The key in the mitochondria segmentation file, analogous to --crista_key."
365    )
366    parser.add_argument(
367        "--voxel_size", type=float,
368        help="The voxel size in nm, applied to all inputs. If not given it is read from the raw tomogram "
369        "passed via --tomogram_path."
370    )
371    parser.add_argument(
372        "--tomogram_path",
373        help="The filepath to the raw tomogram (mrc/rec), or a directory of them, used to read the voxel size "
374        "when --voxel_size is not given."
375    )
376    parser.add_argument(
377        "--membrane_thickness", type=float, default=8.0,
378        help="The membrane shell thickness in nm. By default 8.0."
379    )
380    parser.add_argument(
381        "--border_gap", type=float, default=None,
382        help="The distance from the volume faces where the membrane is suppressed, in nm. "
383        "By default the same as the membrane thickness."
384    )
385    parser.add_argument(
386        "--method", default="skip", choices=["skip", "fast", "exact"],
387        help="How the crista orientation anisotropy is computed. 'skip' (default) does not compute it, "
388        "'fast' uses a downsampled crop (relative only), 'exact' uses the full-resolution structure tensor."
389    )
390    parser.add_argument(
391        "--membrane_mode", default="slice_2d", choices=["slice_2d", "shell_3d"],
392        help="How the membrane shell is built - 'slice_2d' (default, per-Z-slice) or 'shell_3d' (connected 3D shell)."
393    )
394    parser.add_argument(
395        "--n_jobs", type=int, default=-1,
396        help="The number of workers for the per-mitochondrion computation. By default -1 (all cores)."
397    )
398    parser.add_argument(
399        "--force", action="store_true",
400        help="Whether to over-write already present result tables."
401    )
402    parser.add_argument(
403        "--verbose", "-v", action="store_true",
404        help="Whether to show a progress bar over the mitochondria of each file."
405    )
406    args = parser.parse_args()
407
408    cristae_analysis_helper(
409        args.crista_path, args.mito_path, args.output_path,
410        crista_key=args.crista_key, mito_key=args.mito_key,
411        voxel_size=args.voxel_size, tomogram_path=args.tomogram_path,
412        membrane_thickness_nm=args.membrane_thickness, border_gap_nm=args.border_gap,
413        method=args.method, membrane_mode=args.membrane_mode,
414        n_jobs=args.n_jobs, force=args.force, verbose=args.verbose,
415    )