synapse_net.tools.cli

  1import argparse
  2from functools import partial
  3
  4import torch
  5from ..imod.to_imod import export_helper, write_segmentation_to_imod_as_points, write_segmentation_to_imod
  6from ..inference.inference import _get_model_registry, get_model, get_model_training_resolution, run_segmentation
  7from ..inference.util import inference_helper, parse_tiling
  8
  9
 10def imod_point_cli():
 11    parser = argparse.ArgumentParser(
 12        description="Convert a vesicle segmentation to an IMOD point model, "
 13        "corresponding to a sphere for each vesicle in the segmentation."
 14    )
 15    parser.add_argument(
 16        "--input_path", "-i", required=True,
 17        help="The filepath to the mrc file or the directory containing the tomogram data."
 18    )
 19    parser.add_argument(
 20        "--segmentation_path", "-s", required=True,
 21        help="The filepath to the file or the directory containing the segmentations."
 22    )
 23    parser.add_argument(
 24        "--output_path", "-o", required=True,
 25        help="The filepath to directory where the segmentations will be saved."
 26    )
 27    parser.add_argument(
 28        "--segmentation_key", "-k",
 29        help="The key in the segmentation files. If not given we assume that the segmentations are stored as tif."
 30        "If given, we assume they are stored as hdf5 files, and use the key to load the internal dataset."
 31    )
 32    parser.add_argument(
 33        "--min_radius", type=float, default=10.0,
 34        help="The minimum vesicle radius in nm. Objects that are smaller than this radius will be exclded from the export."  # noqa
 35    )
 36    parser.add_argument(
 37        "--radius_factor", type=float, default=1.0,
 38        help="A factor for scaling the sphere radius for the export. "
 39        "This can be used to fit the size of segmented vesicles to the best matching spheres.",
 40    )
 41    parser.add_argument(
 42        "--force", action="store_true",
 43        help="Whether to over-write already present export results."
 44    )
 45    args = parser.parse_args()
 46
 47    export_function = partial(
 48        write_segmentation_to_imod_as_points,
 49        min_radius=args.min_radius,
 50        radius_factor=args.radius_factor,
 51    )
 52
 53    export_helper(
 54        input_path=args.input_path,
 55        segmentation_path=args.segmentation_path,
 56        output_root=args.output_path,
 57        export_function=export_function,
 58        force=args.force,
 59        segmentation_key=args.segmentation_key,
 60    )
 61
 62
 63def imod_object_cli():
 64    parser = argparse.ArgumentParser(
 65        description="Convert segmented objects to close contour IMOD models."
 66    )
 67    parser.add_argument(
 68        "--input_path", "-i", required=True,
 69        help="The filepath to the mrc file or the directory containing the tomogram data."
 70    )
 71    parser.add_argument(
 72        "--segmentation_path", "-s", required=True,
 73        help="The filepath to the file or the directory containing the segmentations."
 74    )
 75    parser.add_argument(
 76        "--output_path", "-o", required=True,
 77        help="The filepath to directory where the segmentations will be saved."
 78    )
 79    parser.add_argument(
 80        "--segmentation_key", "-k",
 81        help="The key in the segmentation files. If not given we assume that the segmentations are stored as tif."
 82        "If given, we assume they are stored as hdf5 files, and use the key to load the internal dataset."
 83    )
 84    parser.add_argument(
 85        "--force", action="store_true",
 86        help="Whether to over-write already present export results."
 87    )
 88    args = parser.parse_args()
 89    export_helper(
 90        input_path=args.input_path,
 91        segmentation_path=args.segmentation_path,
 92        output_root=args.output_path,
 93        export_function=write_segmentation_to_imod,
 94        force=args.force,
 95        segmentation_key=args.segmentation_key,
 96    )
 97
 98
 99# TODO: handle kwargs
100def segmentation_cli():
101    parser = argparse.ArgumentParser(description="Run segmentation.")
102    parser.add_argument(
103        "--input_path", "-i", required=True,
104        help="The filepath to the mrc file or the directory containing the tomogram data."
105    )
106    parser.add_argument(
107        "--output_path", "-o", required=True,
108        help="The filepath to directory where the segmentations will be saved."
109    )
110    model_names = list(_get_model_registry().urls.keys())
111    model_names = ", ".join(model_names)
112    parser.add_argument(
113        "--model", "-m", required=True,
114        help=f"The model type. The following models are currently available: {model_names}"
115    )
116    parser.add_argument(
117        "--mask_path", help="The filepath to a tif file with a mask that will be used to restrict the segmentation."
118        "Can also be a directory with tifs if the filestructure matches input_path."
119    )
120    parser.add_argument("--input_key", "-k", required=False)
121    parser.add_argument(
122        "--force", action="store_true",
123        help="Whether to over-write already present segmentation results."
124    )
125    parser.add_argument(
126        "--tile_shape", type=int, nargs=3,
127        help="The tile shape for prediction, in ZYX order. Lower the tile shape if GPU memory is insufficient."
128    )
129    parser.add_argument(
130        "--halo", type=int, nargs=3,
131        help="The halo for prediction, in ZYX order. Increase the halo to minimize boundary artifacts."
132    )
133    parser.add_argument(
134        "--data_ext", default=".mrc", help="The extension of the tomogram data. By default .mrc."
135    )
136    parser.add_argument(
137        "--checkpoint", "-c", help="Path to a custom model, e.g. from domain adaptation.",
138    )
139    parser.add_argument(
140        "--segmentation_key", "-s",
141        help="If given, the outputs will be saved to an hdf5 file with this key. Otherwise they will be saved as tif.",
142    )
143    parser.add_argument(
144        "--scale", type=float,
145        help="The factor for rescaling the data before inference. "
146        "By default, the scaling factor will be derived from the voxel size of the input data. "
147        "If this parameter is given it will over-ride the default behavior. "
148    )
149    args = parser.parse_args()
150
151    if args.checkpoint is None:
152        model = get_model(args.model)
153    else:
154        model = torch.load(args.checkpoint, weights_only=False)
155        assert model is not None, f"The model from {args.checkpoint} could not be loaded."
156
157    is_2d = "2d" in args.model
158    tiling = parse_tiling(args.tile_shape, args.halo, is_2d=is_2d)
159
160    # If the scale argument is not passed, then we get the average training resolution for the model.
161    # The inputs will then be scaled to match this resolution based on the voxel size from the mrc files.
162    if args.scale is None:
163        model_resolution = get_model_training_resolution(args.model)
164        model_resolution = tuple(model_resolution[ax] for ax in ("yx" if is_2d else "zyx"))
165        scale = None
166    # Otherwise, we set the model resolution to None and use the scaling factor provided by the user.
167    else:
168        model_resolution = None
169        scale = (2 if is_2d else 3) * (args.scale,)
170
171    segmentation_function = partial(
172        run_segmentation, model=model, model_type=args.model, verbose=False, tiling=tiling,
173    )
174    inference_helper(
175        args.input_path, args.output_path, segmentation_function,
176        mask_input_path=args.mask_path, force=args.force, data_ext=args.data_ext,
177        output_key=args.segmentation_key, model_resolution=model_resolution, scale=scale,
178    )