✨ Define SemanticSegmentor with the New EngineABC

requirements/requirements.txt

@@ -4,6 +4,7 @@ aiohttp>=3.8.1
 albumentations>=1.3.0
 bokeh>=3.1.1, <3.6.0
 Click>=8.1.3, <8.2.0
+dask>=0.12.1
 defusedxml>=0.7.1
 filelock>=3.9.0
 flask>=2.2.2

tests/conftest.py

@@ -531,6 +531,7 @@ def sample_wsi_dict(remote_sample: Callable) -> dict:
         "wsi4_4k_4k_svs",
         "wsi3_20k_20k_pred",
         "wsi4_4k_4k_pred",
+        "wsi4_1k_1k_svs",
     ]
     return {name: remote_sample(name) for name in file_names}
 

tests/engines/test_engine_abc.py

@@ -6,10 +6,11 @@
 import logging
 import shutil
 from pathlib import Path
-from typing import TYPE_CHECKING, NoReturn
+from typing import NoReturn
 
 import numpy as np
 import pytest
+import torch
 import torchvision.models as torch_models
 from typing_extensions import Unpack
 
@@ -26,8 +27,7 @@
 )
 from tiatoolbox.models.engine.io_config import ModelIOConfigABC
 
-if TYPE_CHECKING:
-    import torch.nn
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 
 class TestEngineABC(EngineABC):
@@ -69,6 +69,8 @@ def post_process_wsi(
         """Post process WSI output."""
         return super().post_process_wsi(
             raw_predictions=raw_predictions,
+            prediction_shape=(self.batch_size, 1),
+            prediction_dtype=int,
             **kwargs,
         )
 
@@ -79,7 +81,7 @@ def infer_wsi(
         **kwargs: dict,
     ) -> dict | np.ndarray:
         """Test infer_wsi."""
-        return super().infer_wsi(
+        return super().infer_wsi(  # skipcq: PYL-E1121
             dataloader,
             save_path,
             **kwargs,
@@ -362,7 +364,7 @@ def test_engine_run_with_verbose() -> NoReturn:
     out = eng.run(
         images=np.zeros((10, 224, 224, 3), dtype=np.uint8),
         labels=list(range(10)),
-        on_gpu=False,
+        device=device,
     )
 
     assert "probabilities" in out
@@ -457,37 +459,6 @@ def test_patch_pred_zarr_store(tmp_path: pytest.TempPathFactory) -> NoReturn:
         )
 
 
-def test_cache_mode_patches(tmp_path: pytest.TempPathFactory) -> NoReturn:
-    """Test the caching mode."""
-    save_dir = tmp_path / "patch_output"
-
-    eng = TestEngineABC(model="alexnet-kather100k")
-    out = eng.run(
-        images=np.zeros((10, 224, 224, 3), dtype=np.uint8),
-        on_gpu=False,
-        save_dir=save_dir,
-        overwrite=True,
-        cache_mode=True,
-    )
-    assert out.exists(), "Zarr output file does not exist"
-
-    output_file_name = "output2.zarr"
-    cache_size = 4
-    out = eng.run(
-        images=np.zeros((10, 224, 224, 3), dtype=np.uint8),
-        on_gpu=False,
-        save_dir=save_dir,
-        overwrite=True,
-        cache_mode=True,
-        cache_size=4,
-        batch_size=8,
-        output_file=output_file_name,
-    )
-    assert out.stem == output_file_name.split(".")[0]
-    assert eng.batch_size == cache_size
-    assert out.exists(), "Zarr output file does not exist"
-
-
 def test_get_dataloader(sample_svs: Path) -> None:
     """Test the get_dataloader function."""
     eng = TestEngineABC(model="alexnet-kather100k")
@@ -521,7 +492,7 @@ def test_io_config_delegation(tmp_path: Path, caplog: pytest.LogCaptureFixture)
     eng = TestEngineABC(model=model)
 
     kwargs = {
-        "patch_input_shape": [512, 512],
+        "patch_input_shape": [224, 224],
         "input_resolutions": [{"units": "mpp", "resolution": 1.75}],
     }
     with caplog.at_level(logging.WARNING):
@@ -537,7 +508,7 @@ def test_io_config_delegation(tmp_path: Path, caplog: pytest.LogCaptureFixture)
 
     # test providing config / full input info for non pretrained models
     ioconfig = ModelIOConfigABC(
-        patch_input_shape=(512, 512),
+        patch_input_shape=(224, 224),
         stride_shape=(256, 256),
         input_resolutions=[{"resolution": 1.35, "units": "mpp"}],
     )
@@ -547,7 +518,7 @@ def test_io_config_delegation(tmp_path: Path, caplog: pytest.LogCaptureFixture)
         save_dir=f"{tmp_path}/dump",
         ioconfig=ioconfig,
     )
-    assert eng._ioconfig.patch_input_shape == (512, 512)
+    assert eng._ioconfig.patch_input_shape == (224, 224)
     assert eng._ioconfig.stride_shape == (256, 256)
     assert eng._ioconfig.input_resolutions == [{"resolution": 1.35, "units": "mpp"}]
     shutil.rmtree(tmp_path / "dump", ignore_errors=True)
@@ -558,15 +529,15 @@ def test_io_config_delegation(tmp_path: Path, caplog: pytest.LogCaptureFixture)
         save_dir=f"{tmp_path}/dump",
         **kwargs,
     )
-    assert eng._ioconfig.patch_input_shape == [512, 512]
-    assert eng._ioconfig.stride_shape == [512, 512]
+    assert eng._ioconfig.patch_input_shape == [224, 224]
+    assert eng._ioconfig.stride_shape == [224, 224]
     assert eng._ioconfig.input_resolutions == [{"resolution": 1.75, "units": "mpp"}]
     shutil.rmtree(tmp_path / "dump", ignore_errors=True)
 
     # test overwriting pretrained ioconfig
     eng = TestEngineABC(model="alexnet-kather100k")
     eng.run(
-        images=np.zeros((10, 224, 224, 3), dtype=np.uint8),
+        images=np.zeros((10, 300, 300, 3), dtype=np.uint8),
         patch_input_shape=(300, 300),
         stride_shape=(300, 300),
         input_resolutions=[{"units": "baseline", "resolution": 1.99}],
@@ -580,7 +551,7 @@ def test_io_config_delegation(tmp_path: Path, caplog: pytest.LogCaptureFixture)
     shutil.rmtree(tmp_path / "dump", ignore_errors=True)
 
     eng.run(
-        images=np.zeros((10, 224, 224, 3), dtype=np.uint8),
+        images=np.zeros((10, 300, 300, 3), dtype=np.uint8),
         patch_input_shape=(300, 300),
         stride_shape=(300, 300),
         input_resolutions=None,
@@ -618,3 +589,11 @@ def test_io_config_delegation(tmp_path: Path, caplog: pytest.LogCaptureFixture)
                 stride_shape=(1, 1),
                 input_resolutions=_kwargs["input_resolutions"],
             )
+
+
+def test_save_predictions_incorrect_output_type() -> None:
+    """Engine should raise TypeError if incorrect output type is requested."""
+    eng = TestEngineABC(model="alexnet-kather100k")
+
+    with pytest.raises(TypeError, match=r".*Unsupported output type.* "):
+        eng.save_predictions({"predictions": np.zeros((20, 9))}, output_type="random")

tests/engines/test_patch_predictor.py

@@ -54,6 +54,7 @@ def _test_predictor_output(
         cache_mode=cache_mode,
         save_dir=save_dir,
         output_type=output_type,
+        return_probabilities=True,
     )
 
     if tmp_path is not None:
@@ -188,6 +189,7 @@ def test_patch_predictor_api(
     output = predictor.run(
         inputs,
         device="cpu",
+        return_probabilities=True,
     )
     assert sorted(output.keys()) == ["predictions", "probabilities"]
     assert len(output["probabilities"]) == 2
@@ -198,10 +200,11 @@ def test_patch_predictor_api(
         inputs,
         labels=["1", "a"],
         return_labels=True,
+        return_probabilities=True,
     )
     assert sorted(output.keys()) == sorted(["labels", "predictions", "probabilities"])
     assert len(output["probabilities"]) == len(output["labels"])
-    assert output["labels"].tolist() == ["1", "a"]
+    assert list(output["labels"]) == ["1", "a"]
     shutil.rmtree(save_dir_path, ignore_errors=True)
 
     # test loading user weight
@@ -234,10 +237,12 @@ def test_patch_predictor_api(
         labels=[1, 2],
         return_labels=True,
         ioconfig=ioconfig,
+        return_probabilities=True,
+        num_workers=1,
     )
     assert sorted(output.keys()) == sorted(["labels", "predictions", "probabilities"])
     assert len(output["probabilities"]) == len(output["labels"])
-    assert output["labels"].tolist() == [1, 2]
+    assert list(output["labels"]) == [1, 2]
 
 
 def test_wsi_predictor_api(
@@ -274,6 +279,7 @@ def test_wsi_predictor_api(
         images=[mini_wsi_svs, mini_wsi_jpg],
         masks=[mini_wsi_msk, mini_wsi_msk],
         patch_mode=False,
+        return_probabilities=True,
         **_kwargs,
     )
 
@@ -324,7 +330,6 @@ def test_patch_predictor_kather100k_output(
             classification_check=[6, 3],
         )
 
-    # cache mode
     for model, expected_prob in pretrained_info.items():
         _test_predictor_output(
             inputs,