Expose scalable segmentation to the CLI

Author: constantinpape

synapse_net/inference/scalable_segmentation.py

@@ -78,6 +78,7 @@ def scalable_segmentation(
     min_size: int = 500,
     prediction: Optional[ArrayLike] = None,
     verbose: bool = True,
+    mask: Optional[ArrayLike] = None,
 ) -> None:
     """Run segmentation based on a prediction with foreground and boundary channel.
 
@@ -100,6 +101,8 @@ def scalable_segmentation(
             If not given will be stored in a temporary n5 array.
         verbose: Whether to print timing information.
     """
+    if mask is not None:
+        raise NotImplementedError
     assert model.out_channels == 2
 
     # Create a temporary directory for storing the predictions.

synapse_net/inference/util.py

@@ -344,6 +344,7 @@ def inference_helper(
     output_key: Optional[str] = None,
     model_resolution: Optional[Tuple[float, float, float]] = None,
     scale: Optional[Tuple[float, float, float]] = None,
+    allocate_output: bool = False,
 ) -> None:
     """Helper function to run segmentation for mrc files.
 
@@ -366,6 +367,7 @@ def inference_helper(
         model_resolution: The resolution / voxel size to which the inputs should be scaled for prediction.
             If given, the scaling factor will automatically be determined based on the voxel_size of the input data.
         scale: Fixed factor for scaling the model inputs. Cannot be passed together with 'model_resolution'.
+        allocate_output: Whether to allocate the output for the segmentation function.
     """
     if (scale is not None) and (model_resolution is not None):
         raise ValueError("You must not provide both 'scale' and 'model_resolution' arguments.")
@@ -431,7 +433,11 @@ def inference_helper(
             this_scale = _derive_scale(img_path, model_resolution)
 
         # Run the segmentation.
-        segmentation = segmentation_function(input_volume, mask=mask, scale=this_scale)
+        if allocate_output:
+            segmentation = np.zeros(input_volume.shape, dtype="uint32")
+            segmentation_function(input_volume, output=segmentation, mask=mask, scale=this_scale)
+        else:
+            segmentation = segmentation_function(input_volume, mask=mask, scale=this_scale)
 
         # Write the result to tif or h5.
         os.makedirs(os.path.split(output_path)[0], exist_ok=True)

synapse_net/tools/cli.py

@@ -6,6 +6,7 @@
 import torch_em
 from ..imod.to_imod import export_helper, write_segmentation_to_imod_as_points, write_segmentation_to_imod
 from ..inference.inference import _get_model_registry, get_model, get_model_training_resolution, run_segmentation
+from ..inference.scalable_segmentation import scalable_segmentation
 from ..inference.util import inference_helper, parse_tiling
 
 
@@ -152,10 +153,9 @@ def segmentation_cli():
         "--verbose", "-v", action="store_true",
         help="Whether to print verbose information about the segmentation progress."
     )
-    # TODO scalable seg
     parser.add_argument(
-        "--", action="store_true",
-        help=""
+        "--scalable", action="store_true", help="Use the scalable segmentation implementation. "
+        "Currently this only works for vesicles, mitochondria, or active zones."
     )
     args = parser.parse_args()
 
@@ -186,11 +186,26 @@ def segmentation_cli():
         model_resolution = None
         scale = (2 if is_2d else 3) * (args.scale,)
 
-    segmentation_function = partial(
-        run_segmentation, model=model, model_type=args.model, verbose=args.verbose, tiling=tiling,
-    )
+    if args.scalable:
+        if not args.model.startswith(("vesicle", "mito", "active")):
+            raise ValueError(
+                "The scalable segmentation implementation is currently only supported for "
+                f"vesicles, mitochondria, or active zones, not for {args.model}."
+            )
+        segmentation_function = partial(
+            scalable_segmentation, model=model, tiling=tiling, verbose=args.verbose
+        )
+        allocate_output = True
+
+    else:
+        segmentation_function = partial(
+            run_segmentation, model=model, model_type=args.model, verbose=args.verbose, tiling=tiling,
+        )
+        allocate_output = False
+
     inference_helper(
         args.input_path, args.output_path, segmentation_function,
         mask_input_path=args.mask_path, force=args.force, data_ext=args.data_ext,
         output_key=args.segmentation_key, model_resolution=model_resolution, scale=scale,
+        allocate_output=allocate_output
     )