✨ Define SemanticSegmentor with the New EngineABC (#866)

## Summary of Changes

### Major Additions
- **Dask Integration:**  
  - Added `dask` as a dependency and integrated Dask arrays and lazy computation throughout the engine and patch predictor code.
  - Added Dask-based merging, chunking, and memory-aware processing for large images and WSIs.

- **Zarr Output Support:**  
  - Added support for saving model predictions and intermediate results directly to Zarr format.
  - New CLI options and internal logic for Zarr output, including memory thresholding and chunked writes.

- **SemanticSegmentor Engine:**  
  - Added a new `SemanticSegmentor` engine with Dask/Zarr support and new test coverage (`test_semantic_segmentor.py`).
  - Added CLI entrypoint for `semantic_segmentor` and removed the old `semantic_segment` CLI.

- **Enhanced CLI and Config:**  
  - Added CLI options for memory threshold, unified worker options, and improved mask handling.
  - Updated YAML configs and sample data for new models and test images.

- **Utilities and Validation:**  
  - Added utility functions for minimal dtype casting, patch/stride validation, and improved error handling (e.g., `DimensionMismatchError`).
  - Improved annotation store conversion for Dask arrays and Zarr-backed outputs.

- **Changes to `kwarg`**
  - Add `memory-threshold`
  - Unified `num-loader-workers` and `num-postproc-workers` into `num-workers`
  - Removed `cache_mode` as cache mode is automatically handled.

---

### Major Removals/Refactors
- **Removed Old CLI and Redundant Code:**  
  - Deleted the old `semantic_segment.py` CLI and replaced it with `semantic_segmentor.py`.
  - Removed legacy cache mode and patch prediction Zarr store tests.

- **Refactored Model and Dataset APIs:**  
  - Unified and simplified model inference APIs to always return arrays (not dicts) for batch outputs.
  - Refactored dataset classes to enforce patch shape validation and remove legacy “mode” logic.

- **Test Cleanup:**  
  - Removed or updated tests that relied on old APIs or cache mode.
  - Refactored test assertions for new output types and Dask array handling.

- **API Consistency:**  
  - Standardized function and argument names across engines, CLI, and utility modules.
  - Updated docstrings and type hints for clarity and consistency.

---

### Notable File Changes
- **New:**  
  - `tiatoolbox/cli/semantic_segmentor.py`
  - `tests/engines/test_semantic_segmentor.py`

- **Removed:**  
  - `tiatoolbox/cli/semantic_segment.py`
  - Old cache mode and patch Zarr store tests

- **Heavily Modified:**  
  - `engine_abc.py`, `patch_predictor.py`, `semantic_segmentor.py`
  - CLI modules and test suites
  - Dataset and utility modules for Dask/Zarr compatibility

---

### Impact

- Enables scalable, parallel, and memory-efficient inference and output saving for large images.
- Simplifies downstream analysis by supporting Zarr as a native output format.
- Lays the groundwork for further Dask-based optimizations in TIAToolbox.


---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: shaneahmed

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>=2025.10.0
 defusedxml>=0.7.1
 filelock>=3.9.0
 flask>=2.2.2

tests/conftest.py

@@ -534,6 +534,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,
@@ -169,26 +171,26 @@ def test_ioconfig() -> NoReturn:
 
 
 def test_prepare_engines_save_dir(
-    tmp_path: pytest.TempPathFactory,
+    track_tmp_path: pytest.TempPathFactory,
     caplog: pytest.LogCaptureFixture,
 ) -> NoReturn:
     """Test prepare save directory for engines."""
     out_dir = prepare_engines_save_dir(
-        save_dir=tmp_path / "patch_output",
+        save_dir=track_tmp_path / "patch_output",
         patch_mode=True,
         overwrite=False,
     )
 
-    assert out_dir == tmp_path / "patch_output"
+    assert out_dir == track_tmp_path / "patch_output"
     assert out_dir.exists()
 
     out_dir = prepare_engines_save_dir(
-        save_dir=tmp_path / "patch_output",
+        save_dir=track_tmp_path / "patch_output",
         patch_mode=True,
         overwrite=True,
     )
 
-    assert out_dir == tmp_path / "patch_output"
+    assert out_dir == track_tmp_path / "patch_output"
     assert out_dir.exists()
 
     out_dir = prepare_engines_save_dir(
@@ -209,43 +211,43 @@ def test_prepare_engines_save_dir(
         )
 
     out_dir = prepare_engines_save_dir(
-        save_dir=tmp_path / "wsi_single_output",
+        save_dir=track_tmp_path / "wsi_single_output",
         patch_mode=False,
         overwrite=False,
     )
 
-    assert out_dir == tmp_path / "wsi_single_output"
+    assert out_dir == track_tmp_path / "wsi_single_output"
     assert out_dir.exists()
     assert r"When providing multiple whole-slide images / tiles" not in caplog.text
 
     out_dir = prepare_engines_save_dir(
-        save_dir=tmp_path / "wsi_multiple_output",
+        save_dir=track_tmp_path / "wsi_multiple_output",
         patch_mode=False,
         overwrite=False,
     )
 
-    assert out_dir == tmp_path / "wsi_multiple_output"
+    assert out_dir == track_tmp_path / "wsi_multiple_output"
     assert out_dir.exists()
     assert r"When providing multiple whole slide images" in caplog.text
 
     # test for file overwrite with Path.mkdirs() method
     out_path = prepare_engines_save_dir(
-        save_dir=tmp_path / "patch_output" / "output.zarr",
+        save_dir=track_tmp_path / "patch_output" / "output.zarr",
         patch_mode=True,
         overwrite=True,
     )
     assert out_path.exists()
 
     out_path = prepare_engines_save_dir(
-        save_dir=tmp_path / "patch_output" / "output.zarr",
+        save_dir=track_tmp_path / "patch_output" / "output.zarr",
         patch_mode=True,
         overwrite=True,
     )
     assert out_path.exists()
 
     with pytest.raises(FileExistsError):
         out_path = prepare_engines_save_dir(
-            save_dir=tmp_path / "patch_output" / "output.zarr",
+            save_dir=track_tmp_path / "patch_output" / "output.zarr",
             patch_mode=True,
             overwrite=False,
         )
@@ -362,16 +364,16 @@ 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
     assert "labels" in out
 
 
-def test_patch_pred_zarr_store(tmp_path: pytest.TempPathFactory) -> NoReturn:
+def test_patch_pred_zarr_store(track_tmp_path: pytest.TempPathFactory) -> NoReturn:
     """Test the engine run and patch pred store."""
-    save_dir = tmp_path / "patch_output"
+    save_dir = track_tmp_path / "patch_output"
 
     eng = TestEngineABC(model="alexnet-kather100k")
     out = eng.run(
@@ -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")
@@ -514,82 +485,84 @@ def test_get_dataloader(sample_svs: Path) -> None:
     assert isinstance(dataloader.dataset, WSIPatchDataset)
 
 
-def test_io_config_delegation(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
+def test_io_config_delegation(
+    track_tmp_path: Path, caplog: pytest.LogCaptureFixture
+) -> None:
     """Test for delegating args to io config."""
     # test not providing config / full input info for not pretrained models
     model = CNNModel("resnet50")
     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):
         eng.run(
             np.zeros((10, 224, 224, 3)),
             patch_mode=True,
-            save_dir=tmp_path / "dump",
+            save_dir=track_tmp_path / "dump",
             patch_input_shape=kwargs["patch_input_shape"],
             input_resolutions=kwargs["input_resolutions"],
         )
         assert "provide a valid ModelIOConfigABC" in caplog.text
-    shutil.rmtree(tmp_path / "dump", ignore_errors=True)
+    shutil.rmtree(track_tmp_path / "dump", ignore_errors=True)
 
     # 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"}],
     )
     eng.run(
         images=np.zeros((10, 224, 224, 3), dtype=np.uint8),
         patch_mode=True,
-        save_dir=f"{tmp_path}/dump",
+        save_dir=f"{track_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)
+    shutil.rmtree(track_tmp_path / "dump", ignore_errors=True)
 
     eng.run(
         images=np.zeros((10, 224, 224, 3), dtype=np.uint8),
         patch_mode=True,
-        save_dir=f"{tmp_path}/dump",
+        save_dir=f"{track_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)
+    shutil.rmtree(track_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}],
         patch_mode=True,
-        save_dir=f"{tmp_path}/dump",
+        save_dir=f"{track_tmp_path}/dump",
     )
     assert eng._ioconfig.patch_input_shape == (300, 300)
     assert eng._ioconfig.stride_shape == (300, 300)
     assert eng._ioconfig.input_resolutions[0]["resolution"] == 1.99
     assert eng._ioconfig.input_resolutions[0]["units"] == "baseline"
-    shutil.rmtree(tmp_path / "dump", ignore_errors=True)
+    shutil.rmtree(track_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,
         patch_mode=True,
-        save_dir=f"{tmp_path}/dump",
+        save_dir=f"{track_tmp_path}/dump",
     )
     assert eng._ioconfig.patch_input_shape == (300, 300)
     assert eng._ioconfig.stride_shape == (300, 300)
-    shutil.rmtree(tmp_path / "dump", ignore_errors=True)
+    shutil.rmtree(track_tmp_path / "dump", ignore_errors=True)
 
     eng.ioconfig = None
     _ioconfig = eng._update_ioconfig(
@@ -618,3 +591,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")