🔥 Remove cache_mode

Author: shaneahmed

tests/engines/test_engine_abc.py

@@ -459,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")

tests/engines/test_semantic_segmentor.py

@@ -37,7 +37,7 @@ def test_semantic_segmentor_patches(remote_sample: Callable, tmp_path: Path) ->
 
     inputs = [sample_image, sample_image]
 
-    assert segmentor.cache_mode is False
+    assert not segmentor.patch_mode
 
     output = segmentor.run(
         images=inputs,

tiatoolbox/models/engine/engine_abc.py

@@ -129,10 +129,6 @@ class EngineABCRunParams(TypedDict, total=False):
     Optional Keys:
         batch_size (int):
             Number of image patches per forward pass.
-        cache_mode (bool):
-            Whether to use caching for large datasets.
-        cache_size (int):
-            Number of patches to process in a batch when caching.
         class_dict (dict):
             Mapping of classification outputs to class names.
         device (str):
@@ -173,8 +169,6 @@ class EngineABCRunParams(TypedDict, total=False):
     """
 
     batch_size: int
-    cache_mode: bool
-    cache_size: int
     class_dict: dict
     device: str
     ioconfig: ModelIOConfigABC
@@ -277,10 +271,6 @@ class EngineABC(ABC):  # noqa: B024
             `stride_shape=patch_input_shape`.
         batch_size (int):
             Number of images fed into the model each time.
-        cache_mode (bool):
-            Whether to use caching for large datasets.
-        cache_size (int):
-            Number of patches to process in a batch when caching.
         labels (list | None):
             Optional labels for input images.
             Only a single label per image is supported.
@@ -370,8 +360,6 @@ def __init__(
         )
         self._ioconfig = self.ioconfig  # runtime ioconfig
         self.batch_size = batch_size
-        self.cache_mode: bool = False
-        self.cache_size: int = self.batch_size if self.batch_size else 10000
         self.labels: list | None = None
         self.num_loader_workers = num_loader_workers
         self.num_post_proc_workers = num_post_proc_workers
@@ -1065,22 +1053,13 @@ def _update_run_params(
             self.drop_keys.append("label")
 
         self.patch_mode = patch_mode
-        if not self.patch_mode:
-            self.cache_mode = True  # if input is WSI run using cache mode.
-
-        if self.cache_mode and self.batch_size > self.cache_size:
-            self.batch_size = self.cache_size
 
         self._validate_input_numbers(images=images, masks=masks, labels=labels)
         if output_type.lower() not in ["dict", "zarr", "annotationstore"]:
             msg = "output_type must be 'dict' or 'zarr' or 'annotationstore'."
             raise TypeError(msg)
 
         self.output_type = output_type
-        if self.cache_mode and output_type.lower() not in ["zarr", "annotationstore"]:
-            self.output_type = "zarr"
-            msg = "output_type has been updated to 'zarr' for cache_mode=True."
-            logger.info(msg)
 
         if save_dir is not None and output_type.lower() not in [
             "zarr",
@@ -1149,7 +1128,7 @@ def _run_patch_mode(
 
         """
         save_path = None
-        if self.cache_mode or save_dir:
+        if save_dir:
             output_file = Path(kwargs.get("output_file", "output.zarr"))
             save_path = save_dir / (str(output_file.stem) + ".zarr")
 

tiatoolbox/models/engine/patch_predictor.py

@@ -4,8 +4,6 @@
 
 from typing import TYPE_CHECKING
 
-import dask.array as da
-from dask import delayed
 from typing_extensions import Unpack
 
 from .engine_abc import EngineABC, EngineABCRunParams
@@ -14,6 +12,7 @@
     import os
     from pathlib import Path
 
+    import dask.array as da
     import numpy as np
 
     from tiatoolbox.annotation import AnnotationStore
@@ -190,10 +189,6 @@ class PatchPredictor(EngineABC):
             at requested read resolution, not with respect to
             level 0, and must be positive. If not provided,
             `stride_shape=patch_input_shape`.
-        cache_mode (bool):
-            Whether to use caching for large datasets.
-        cache_size (int):
-            Number of patches to process in a batch when caching.
         labels (list | None):
             Optional labels for input images.
             Only a single label per image is supported.
@@ -310,12 +305,9 @@ def post_process_patches(
 
         """
         _ = kwargs.get("return_probabilities")
-        raw_predictions = delayed(self.model.postproc_func)(raw_predictions)
-        return da.from_delayed(
-            raw_predictions,
-            shape=prediction_shape,
-            dtype=prediction_dtype,
-        )
+        _ = prediction_shape
+        _ = prediction_dtype
+        return self.model.postproc_func(raw_predictions)
 
     def post_process_wsi(
         self: PatchPredictor,

tiatoolbox/models/engine/semantic_segmentor.py

@@ -82,10 +82,6 @@ class SemanticSegmentorRunParams(PredictorRunParams, total=False):
     Attributes:
         batch_size (int):
             Number of image patches to feed to the model in a forward pass.
-        cache_mode (bool):
-            Whether to run the engine in cache mode. Recommended for large datasets.
-        cache_size (int):
-            Number of patches to process in a batch when cache_mode is True.
         class_dict (dict):
             Optional dictionary mapping classification outputs to class names.
         device (str):
@@ -223,10 +219,6 @@ class SemanticSegmentor(PatchPredictor):
             at requested read resolution, not with respect to
             level 0, and must be positive. If not provided,
             `stride_shape=patch_input_shape`.
-        cache_mode (bool):
-            Whether to use caching for large datasets.
-        cache_size (int):
-            Number of patches to process in a batch when caching.
         labels (list | None):
             Optional labels for input images.
             Only a single label per image is supported.
@@ -587,7 +579,7 @@ def save_predictions(
             )
             save_paths = out_file
 
-        if return_probabilities and self.cache_mode:
+        if return_probabilities:
             zarr_save_path = save_path.parent.with_suffix(".zarr")
             msg = (
                 f"Probability maps cannot be saved as AnnotationStore. "