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