♻️ Use input_resolutions instead of resolution in Engine Params (#917)

* ♻️ Use `input_resolutions` instead of resolution
- Use `input_resolutions` instead of resolution to make engines outputs compatible with ioconfig.
- Uses input resolution as a list of dictionaries on units and resolution.

* ♻️ Use `input_resolutions` instead of resolution
- Use `input_resolutions` instead of resolution to make engines outputs compatible with ioconfig.
- Uses input resolution as a list of dictionaries on units and resolution.

* ✅ Add test to cli.

Author: shaneahmed

tests/engines/test_engine_abc.py

@@ -521,17 +521,15 @@ def test_io_config_delegation(tmp_path: Path, caplog: pytest.LogCaptureFixture)
 
     kwargs = {
         "patch_input_shape": [512, 512],
-        "resolution": 1.75,
-        "units": "mpp",
+        "input_resolutions": [{"units": "mpp", "resolution": 1.75}],
     }
     with caplog.at_level(logging.WARNING):
         eng.run(
             np.zeros((10, 224, 224, 3)),
             patch_mode=True,
             save_dir=tmp_path / "dump",
             patch_input_shape=kwargs["patch_input_shape"],
-            resolution=kwargs["resolution"],
-            units=kwargs["units"],
+            input_resolutions=kwargs["input_resolutions"],
         )
         assert "provide a valid ModelIOConfigABC" in caplog.text
     shutil.rmtree(tmp_path / "dump", ignore_errors=True)
@@ -570,8 +568,7 @@ def test_io_config_delegation(tmp_path: Path, caplog: pytest.LogCaptureFixture)
         images=np.zeros((10, 224, 224, 3), dtype=np.uint8),
         patch_input_shape=(300, 300),
         stride_shape=(300, 300),
-        resolution=1.99,
-        units="baseline",
+        input_resolutions=[{"units": "baseline", "resolution": 1.99}],
         patch_mode=True,
         save_dir=f"{tmp_path}/dump",
     )
@@ -585,8 +582,7 @@ def test_io_config_delegation(tmp_path: Path, caplog: pytest.LogCaptureFixture)
         images=np.zeros((10, 224, 224, 3), dtype=np.uint8),
         patch_input_shape=(300, 300),
         stride_shape=(300, 300),
-        resolution=None,
-        units=None,
+        input_resolutions=None,
         patch_mode=True,
         save_dir=f"{tmp_path}/dump",
     )
@@ -599,8 +595,7 @@ def test_io_config_delegation(tmp_path: Path, caplog: pytest.LogCaptureFixture)
         ioconfig=None,
         patch_input_shape=(300, 300),
         stride_shape=(300, 300),
-        resolution=1.99,
-        units="baseline",
+        input_resolutions=[{"units": "baseline", "resolution": 1.99}],
     )
 
     assert _ioconfig.patch_input_shape == (300, 300)
@@ -614,12 +609,11 @@ def test_io_config_delegation(tmp_path: Path, caplog: pytest.LogCaptureFixture)
         with pytest.raises(
             ValueError,
             match=r".*Must provide either `ioconfig` or "
-            r"`patch_input_shape`, `resolution`, and `units`*",
+            r"`patch_input_shape` and `input_resolutions`*",
         ):
             eng._update_ioconfig(
                 ioconfig=None,
                 patch_input_shape=_kwargs["patch_input_shape"],
                 stride_shape=(1, 1),
-                resolution=_kwargs["resolution"],
-                units=_kwargs["units"],
+                input_resolutions=_kwargs["input_resolutions"],
             )

tests/engines/test_patch_predictor.py

@@ -1,4 +1,4 @@
-"""Test for Patch Predictor."""
+"""Test PatchPredictor."""
 
 from __future__ import annotations
 
@@ -11,12 +11,14 @@
 
 import numpy as np
 import torch
+import yaml
 import zarr
 from click.testing import CliRunner
 
 from tests.conftest import timed
 from tiatoolbox import cli, logger, rcParam
 from tiatoolbox.models import IOPatchPredictorConfig
+from tiatoolbox.models.architecture import fetch_pretrained_weights
 from tiatoolbox.models.architecture.vanilla import CNNModel
 from tiatoolbox.models.engine.patch_predictor import PatchPredictor
 from tiatoolbox.utils import env_detection as toolbox_env
@@ -86,8 +88,7 @@ def test_io_config_delegation(remote_sample: Callable, tmp_path: Path) -> None:
     predictor = PatchPredictor(model=model, weights=None)
     kwargs = {
         "patch_input_shape": [512, 512],
-        "resolution": 1.75,
-        "units": "mpp",
+        "input_resolutions": [{"units": "mpp", "resolution": 1.75}],
     }
 
     # test providing config / full input info for default models without weights
@@ -134,7 +135,7 @@ def test_io_config_delegation(remote_sample: Callable, tmp_path: Path) -> None:
 
     predictor.run(
         images=[mini_wsi_svs],
-        resolution=1.99,
+        input_resolutions=[{"units": "mpp", "resolution": 1.99}],
         patch_mode=False,
         save_dir=f"{tmp_path}/dump",
     )
@@ -143,7 +144,7 @@ def test_io_config_delegation(remote_sample: Callable, tmp_path: Path) -> None:
 
     predictor.run(
         images=[mini_wsi_svs],
-        units="baseline",
+        input_resolutions=[{"units": "baseline", "resolution": 1.0}],
         patch_mode=False,
         save_dir=f"{tmp_path}/dump",
     )
@@ -152,8 +153,7 @@ def test_io_config_delegation(remote_sample: Callable, tmp_path: Path) -> None:
 
     predictor.run(
         images=[mini_wsi_svs],
-        units="level",
-        resolution=0,
+        input_resolutions=[{"units": "level", "resolution": 0}],
         patch_mode=False,
         save_dir=f"{tmp_path}/dump",
     )
@@ -163,8 +163,7 @@ def test_io_config_delegation(remote_sample: Callable, tmp_path: Path) -> None:
 
     predictor.run(
         images=[mini_wsi_svs],
-        units="power",
-        resolution=20,
+        input_resolutions=[{"units": "power", "resolution": 20}],
         patch_mode=False,
         save_dir=f"{tmp_path}/dump",
     )
@@ -262,8 +261,7 @@ def test_wsi_predictor_api(
     kwargs = {
         "patch_input_shape": patch_size,
         "stride_shape": patch_size,
-        "resolution": 1.0,
-        "units": "baseline",
+        "input_resolutions": [{"units": "baseline", "resolution": 1.0}],
         "save_dir": save_dir,
     }
     # ! add this test back once the read at `baseline` is fixed
@@ -646,6 +644,17 @@ def test_cli_model_multiple_file_mask(remote_sample: Callable, tmp_path: Path) -
     dir_path_masks = tmp_path.joinpath("new_copies_masks")
     dir_path_masks.mkdir()
 
+    config = {
+        "input_resolutions": [{"units": "mpp", "resolution": 0.5}],
+        "patch_input_shape": [224, 224],
+    }
+
+    with Path.open(tmp_path.joinpath("config.yaml"), "w") as fptr:
+        yaml.dump(config, fptr)
+
+    model = "alexnet-kather100k"
+    weights = fetch_pretrained_weights(model)
+
     try:
         dir_path.joinpath("1_" + mini_wsi_svs.name).symlink_to(mini_wsi_svs)
         dir_path.joinpath("2_" + mini_wsi_svs.name).symlink_to(mini_wsi_svs)
@@ -675,6 +684,12 @@ def test_cli_model_multiple_file_mask(remote_sample: Callable, tmp_path: Path) -
             str(False),
             "--masks",
             str(dir_path_masks),
+            "--model",
+            model,
+            "--weights",
+            str(weights),
+            "--yaml-config-path",
+            tmp_path / "config.yaml",
             "--output-path",
             str(tmp_path / "output"),
             "--output-type",

tiatoolbox/cli/common.py

@@ -619,17 +619,17 @@ def prepare_model_cli(
 tiatoolbox_cli = TIAToolboxCLI()
 
 
-def prepare_ioconfig_seg(
-    segment_config_class: type[IOConfigABC],
+def prepare_ioconfig(
+    config_class: type[IOConfigABC],
     pretrained_weights: str | Path | None,
     yaml_config_path: str | Path,
 ) -> IOConfigABC | None:
-    """Prepare ioconfig for segmentation."""
+    """Prepare ioconfig for CLI."""
     import yaml
 
     if pretrained_weights is not None:
         with Path(yaml_config_path).open() as registry_handle:
             ioconfig = yaml.safe_load(registry_handle)
-        return segment_config_class(**ioconfig)
+        return config_class(**ioconfig)
 
     return None

tiatoolbox/cli/nucleus_instance_segment.py

@@ -19,7 +19,7 @@
     cli_pretrained_weights,
     cli_verbose,
     cli_yaml_config_path,
-    prepare_ioconfig_seg,
+    prepare_ioconfig,
     prepare_model_cli,
     tiatoolbox_cli,
 )
@@ -77,7 +77,7 @@ def nucleus_instance_segment(
         file_types=file_types,
     )
 
-    ioconfig = prepare_ioconfig_seg(
+    ioconfig = prepare_ioconfig(
         IOInstanceSegmentorConfig,
         pretrained_weights,
         yaml_config_path,

tiatoolbox/cli/patch_predictor.py

@@ -13,12 +13,12 @@
     cli_output_path,
     cli_output_type,
     cli_patch_mode,
-    cli_resolution,
     cli_return_labels,
     cli_return_probabilities,
-    cli_units,
     cli_verbose,
     cli_weights,
+    cli_yaml_config_path,
+    prepare_ioconfig,
     prepare_model_cli,
     tiatoolbox_cli,
 )
@@ -37,8 +37,7 @@
 @cli_weights()
 @cli_device(default="cpu")
 @cli_batch_size(default=1)
-@cli_resolution(default=0.5)
-@cli_units(default="mpp")
+@cli_yaml_config_path()
 @cli_masks(default=None)
 @cli_num_loader_workers(default=0)
 @cli_output_type(
@@ -56,8 +55,7 @@ def patch_predictor(
     masks: str | None,
     output_path: str,
     batch_size: int,
-    resolution: float,
-    units: str,
+    yaml_config_path: str,
     num_loader_workers: int,
     device: str,
     output_type: str,
@@ -68,6 +66,7 @@ def patch_predictor(
     verbose: bool,
 ) -> None:
     """Process an image/directory of input images with a patch classification CNN."""
+    from tiatoolbox.models.engine.io_config import IOPatchPredictorConfig
     from tiatoolbox.models.engine.patch_predictor import PatchPredictor
 
     files_all, masks_all, output_path = prepare_model_cli(
@@ -85,12 +84,17 @@ def patch_predictor(
         verbose=verbose,
     )
 
+    ioconfig = prepare_ioconfig(
+        IOPatchPredictorConfig,
+        pretrained_weights=weights,
+        yaml_config_path=yaml_config_path,
+    )
+
     _ = predictor.run(
         images=files_all,
         masks=masks_all,
         patch_mode=patch_mode,
-        resolution=resolution,
-        units=units,
+        ioconfig=ioconfig,
         device=device,
         save_dir=output_path,
         output_type=output_type,

tiatoolbox/cli/semantic_segment.py

@@ -17,7 +17,7 @@
     cli_pretrained_weights,
     cli_verbose,
     cli_yaml_config_path,
-    prepare_ioconfig_seg,
+    prepare_ioconfig,
     prepare_model_cli,
     tiatoolbox_cli,
 )
@@ -71,7 +71,7 @@ def semantic_segment(
         file_types=file_types,
     )
 
-    ioconfig = prepare_ioconfig_seg(
+    ioconfig = prepare_ioconfig(
         IOSegmentorConfig,
         pretrained_weights,
         yaml_config_path,

tiatoolbox/models/architecture/utils.py

@@ -3,13 +3,17 @@
 from __future__ import annotations
 
 import sys
+from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
 from torch import nn
 
 from tiatoolbox import logger
 
+if TYPE_CHECKING:  # pragma: no cover
+    from tiatoolbox.models.models_abc import ModelABC
+
 
 def is_torch_compile_compatible() -> bool:
     """Check if the current GPU is compatible with torch-compile.
@@ -45,10 +49,10 @@ def is_torch_compile_compatible() -> bool:
 
 
 def compile_model(
-    model: nn.Module | None = None,
+    model: nn.Module | ModelABC | None = None,
     *,
     mode: str = "default",
-) -> nn.Module:
+) -> torch.nn.Module | ModelABC:
     """A decorator to compile a model using torch-compile.
 
     Args:
@@ -67,7 +71,7 @@ def compile_model(
               CUDA graphs
 
     Returns:
-        torch.nn.Module:
+        torch.nn.Module or ModelABC:
             Compiled model.
 
     """