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