✨ Add CLI interface

Author: shaneahmed

tests/engines/test_semantic_segmentor.py

@@ -10,7 +10,9 @@
 import numpy as np
 import torch
 import zarr
+from click.testing import CliRunner
 
+from tiatoolbox import cli
 from tiatoolbox.annotation import SQLiteStore
 from tiatoolbox.models.engine.semantic_segmentor import SemanticSegmentor
 from tiatoolbox.utils import env_detection as toolbox_env
@@ -354,3 +356,28 @@ def test_wsi_segmentor_annotationstore(
     zarr_group = zarr.open(output[sample_svs].with_suffix(".zarr"), mode="r")
     assert "probabilities" in zarr_group
     assert "Probability maps cannot be saved as AnnotationStore." in caplog.text
+
+
+# -------------------------------------------------------------------------------------
+# Command Line Interface
+# -------------------------------------------------------------------------------------
+
+
+def test_cli_model_single_file(sample_svs: Path, tmp_path: Path) -> None:
+    """Test for models CLI single file."""
+    runner = CliRunner()
+    models_wsi_result = runner.invoke(
+        cli.main,
+        [
+            "semantic-segmentor",
+            "--img-input",
+            str(sample_svs),
+            "--patch-mode",
+            "False",
+            "--output-path",
+            str(tmp_path / "output"),
+        ],
+    )
+
+    assert models_wsi_result.exit_code == 0
+    assert (tmp_path / "output" / (sample_svs.stem + ".db")).exists()

tiatoolbox/cli/__init__.py

@@ -11,7 +11,7 @@
 from tiatoolbox.cli.patch_predictor import patch_predictor
 from tiatoolbox.cli.read_bounds import read_bounds
 from tiatoolbox.cli.save_tiles import save_tiles
-from tiatoolbox.cli.semantic_segment import semantic_segment
+from tiatoolbox.cli.semantic_segmentor import semantic_segmentor
 from tiatoolbox.cli.show_wsi import show_wsi
 from tiatoolbox.cli.slide_info import slide_info
 from tiatoolbox.cli.slide_thumbnail import slide_thumbnail
@@ -42,7 +42,7 @@ def main() -> click.BaseCommand:
 main.add_command(patch_predictor)
 main.add_command(read_bounds)
 main.add_command(save_tiles)
-main.add_command(semantic_segment)
+main.add_command(semantic_segmentor)
 main.add_command(slide_info)
 main.add_command(slide_thumbnail)
 main.add_command(tissue_mask)

tiatoolbox/cli/common.py

@@ -387,14 +387,28 @@ def cli_masks(
     )
 
 
-def cli_auto_generate_mask(
+def cli_memory_threshold(
+    usage_help: str = (
+        "Memory usage threshold (in percentage) to trigger caching behavior."
+    ),
+    default: int = 80,
+) -> Callable:
+    """Enables --batch-size option for cli."""
+    return click.option(
+        "--memory-threshold",
+        help=add_default_to_usage_help(usage_help, default=default),
+        default=default,
+    )
+
+
+def cli_auto_get_mask(
     usage_help: str = "Automatically generate tile/WSI tissue mask.",
     *,
     default: bool = False,
 ) -> Callable:
     """Enables --auto-generate-mask option for cli."""
     return click.option(
-        "--auto-generate-mask",
+        "--auto-get-mask",
         help=add_default_to_usage_help(usage_help, default=default),
         type=bool,
         default=default,
@@ -415,27 +429,14 @@ def cli_yaml_config_path(
     )
 
 
-def cli_num_loader_workers(
+def cli_num_workers(
     usage_help: str = "Number of workers to load the data. Please note that they will "
     "also perform preprocessing.",
     default: int = 0,
 ) -> Callable:
     """Enables --num-loader-workers option for cli."""
     return click.option(
-        "--num-loader-workers",
-        help=add_default_to_usage_help(usage_help, default=default),
-        type=int,
-        default=default,
-    )
-
-
-def cli_num_postproc_workers(
-    usage_help: str = "Number of workers to post-process the network output.",
-    default: int = 0,
-) -> Callable:
-    """Enables --num-postproc-workers option for cli."""
-    return click.option(
-        "--num-postproc-workers",
+        "--num-workers",
         help=add_default_to_usage_help(usage_help, default=default),
         type=int,
         default=default,

tiatoolbox/cli/nucleus_instance_segment.py

@@ -5,15 +5,14 @@
 import click
 
 from tiatoolbox.cli.common import (
-    cli_auto_generate_mask,
+    cli_auto_get_mask,
     cli_batch_size,
     cli_device,
     cli_file_type,
     cli_img_input,
     cli_masks,
     cli_mode,
-    cli_num_loader_workers,
-    cli_num_postproc_workers,
+    cli_num_workers,
     cli_output_path,
     cli_pretrained_model,
     cli_pretrained_weights,
@@ -45,10 +44,9 @@
 @cli_batch_size()
 @cli_masks(default=None)
 @cli_yaml_config_path(default=None)
-@cli_num_loader_workers()
+@cli_num_workers()
 @cli_verbose(default=True)
-@cli_num_postproc_workers(default=0)
-@cli_auto_generate_mask(default=False)
+@cli_auto_get_mask(default=False)
 def nucleus_instance_segment(
     pretrained_model: str,
     pretrained_weights: str,
@@ -60,7 +58,6 @@ def nucleus_instance_segment(
     batch_size: int,
     yaml_config_path: str,
     num_loader_workers: int,
-    num_postproc_workers: int,
     device: str,
     *,
     auto_generate_mask: bool,
@@ -91,7 +88,6 @@ def nucleus_instance_segment(
         pretrained_weights=pretrained_weights,
         batch_size=batch_size,
         num_loader_workers=num_loader_workers,
-        num_postproc_workers=num_postproc_workers,
         auto_generate_mask=auto_generate_mask,
         verbose=verbose,
     )

tiatoolbox/cli/patch_predictor.py

@@ -3,13 +3,15 @@
 from __future__ import annotations
 
 from tiatoolbox.cli.common import (
+    cli_auto_get_mask,
     cli_batch_size,
     cli_device,
     cli_file_type,
     cli_img_input,
     cli_masks,
+    cli_memory_threshold,
     cli_model,
-    cli_num_loader_workers,
+    cli_num_workers,
     cli_output_path,
     cli_output_type,
     cli_patch_mode,
@@ -39,13 +41,15 @@
 @cli_batch_size(default=1)
 @cli_yaml_config_path()
 @cli_masks(default=None)
-@cli_num_loader_workers(default=0)
+@cli_num_workers(default=0)
 @cli_output_type(
     default="AnnotationStore",
 )
+@cli_memory_threshold(default=80)
 @cli_patch_mode(default=False)
 @cli_return_probabilities(default=True)
 @cli_return_labels(default=False)
+@cli_auto_get_mask(default=True)
 @cli_verbose(default=True)
 def patch_predictor(
     model: str,
@@ -56,16 +60,18 @@ def patch_predictor(
     output_path: str,
     batch_size: int,
     yaml_config_path: str,
-    num_loader_workers: int,
+    num_workers: int,
     device: str,
     output_type: str,
+    memory_threshold: int,
     *,
+    patch_mode: bool,
     return_probabilities: bool,
     return_labels: bool,
-    patch_mode: bool,
+    auto_get_mask: bool,
     verbose: bool,
 ) -> None:
-    """Process an image/directory of input images with a patch classification CNN."""
+    """Process an image/directory of input images with a patch classification engine."""
     from tiatoolbox.models.engine.io_config import (  # noqa: PLC0415
         IOPatchPredictorConfig,
     )
@@ -82,7 +88,7 @@ def patch_predictor(
         model=model,
         weights=weights,
         batch_size=batch_size,
-        num_workers=num_loader_workers,
+        num_workers=num_workers,
         verbose=verbose,
     )
 
@@ -102,4 +108,6 @@ def patch_predictor(
         output_type=output_type,
         return_probabilities=return_probabilities,
         return_labels=return_labels,
+        auto_get_mask=auto_get_mask,
+        memory_threshold=memory_threshold,
     )