replace cache_dir with temp dir

Author: Jiaqi-Lv

tests/engines/test_nucleus_detection_engine.py

@@ -256,7 +256,6 @@ def test_nucleus_detector_wsi(remote_sample: Callable, track_tmp_path: Path) ->
         memory_threshold=50,
         images=[mini_wsi_svs],
         save_dir=save_dir,
-        cache_dir=save_dir,
         overwrite=True,
         batch_size=8,
         class_dict={0: "test_nucleus"},
@@ -278,7 +277,6 @@ def test_nucleus_detector_wsi(remote_sample: Callable, track_tmp_path: Path) ->
         memory_threshold=50,
         images=[mini_wsi_svs],
         save_dir=save_dir,
-        cache_dir=save_dir,
         overwrite=True,
         batch_size=8,
     )

tiatoolbox/cli/common.py

@@ -81,19 +81,6 @@ def cli_output_path(
     )
 
 
-def cli_cache_path(
-    usage_help: str = "Path to cache directory to save the cache.",
-    default: str | None = None,
-) -> Callable:
-    """Enables --cache-path option for cli."""
-    return click.option(
-        "--cache-path",
-        help=add_default_to_usage_help(usage_help, default=default),
-        type=str,
-        default=default,
-    )
-
-
 def cli_output_file(
     usage_help: str = "Filename for saving output (e.g., '.zarr' or '.db').",
     default: str | None = None,

tiatoolbox/cli/nucleus_detector.py

@@ -2,13 +2,11 @@
 
 from __future__ import annotations
 
-import shutil
 from typing import TYPE_CHECKING
 
 from tiatoolbox.cli.common import (
     cli_auto_get_mask,
     cli_batch_size,
-    cli_cache_path,
     cli_class_dict,
     cli_device,
     cli_file_type,
@@ -51,10 +49,6 @@
     usage_help="Output directory where model prediction will be saved.",
     default="nucleus_detection",
 )
-@cli_cache_path(
-    usage_help="Directory to use for caching intermediate files.",
-    default="nucleus_detection",
-)
 @cli_output_file(default=None)
 @cli_file_type(
     default="*.png, *.jpg, *.jpeg, *.tif, *.tiff, *.svs, *.ndpi, *.jp2, *.mrxs",
@@ -96,7 +90,6 @@ def nucleus_detector(
     output_resolutions: list[dict],
     masks: str | None,
     output_path: str,
-    cache_path: str,
     patch_input_shape: IntPair | None,
     patch_output_shape: tuple[int, int] | None,
     stride_shape: IntPair | None,
@@ -158,7 +151,6 @@ def nucleus_detector(
         ioconfig=ioconfig,
         device=device,
         save_dir=output_path,
-        cache_dir=cache_path,
         output_type=output_type,
         return_probabilities=return_probabilities,
         auto_get_mask=auto_get_mask,
@@ -174,5 +166,3 @@ def nucleus_detector(
         overwrite=overwrite,
         verbose=verbose,
     )
-
-    shutil.rmtree(cache_path)

tiatoolbox/models/engine/nucleus_detector.py

@@ -47,6 +47,7 @@
 from __future__ import annotations
 
 import shutil
+import tempfile
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -116,8 +117,6 @@ class NucleusDetectorRunParams(SemanticSegmentorRunParams, total=False):
         postproc_tile_shape (tuple[int, int]):
             Tile shape (height, width) used during post-processing
             (in pixels) to control rechunking behavior.
-        cache_dir (str or os.PathLike):
-            Directory for caching intermediate results during WSI processing.
         return_labels (bool):
             Whether to return labels with predictions.
         return_probabilities (bool):
@@ -136,7 +135,6 @@ class NucleusDetectorRunParams(SemanticSegmentorRunParams, total=False):
     threshold_abs: float
     threshold_rel: float
     postproc_tile_shape: IntPair
-    cache_dir: str | os.PathLike
 
 
 class NucleusDetector(SemanticSegmentor):
@@ -478,17 +476,18 @@ def post_process_wsi(
             depth_w=depth_w,
         )
 
-        logger.info("Computing and saving centroid maps to cache as zarr.")
-        save_dir = kwargs.get("cache_dir", "./tmp/")
-        save_path = Path(save_dir) / "detection_maps"
-        if save_path.exists():
-            shutil.rmtree(save_path)
-
-        task = centroid_maps.to_zarr(url=save_path, compute=False, object_codec=None)
+        logger.info("Computing and saving centroid maps to temporary zarr file.")
+        temp_zarr_file = tempfile.TemporaryDirectory(
+            prefix="tiatoolbox_nucleus_detector_", suffix=".zarr"
+        )
+        logger.info("Temporary zarr file created at: %s", temp_zarr_file.name)
+        task = centroid_maps.to_zarr(
+            url=temp_zarr_file.name, compute=False, object_codec=None
+        )
         with ProgressBar():
             compute(task)
 
-        centroid_maps = da.from_zarr(save_path)
+        centroid_maps = da.from_zarr(temp_zarr_file.name)
 
         return self._centroid_maps_to_detection_arrays(centroid_maps)
 
@@ -1112,7 +1111,7 @@ class names.
 
 
         """
-        return super().run(
+        output = super().run(
             images=images,
             masks=masks,
             input_resolutions=input_resolutions,
@@ -1124,3 +1123,21 @@ class names.
             output_type=output_type,
             **kwargs,
         )
+
+        if not patch_mode:
+            # Clean up temporary zarr directory after WSI processing
+            # It should have been already deleted, but check anyway
+            temp_dir = Path(tempfile.gettempdir())
+            if temp_dir.exists():
+                # find file starting with 'tiatoolbox_nucleus_detector_'
+                # and ending with '.zarr'
+                for item in temp_dir.iterdir():
+                    if item.name.startswith(
+                        "tiatoolbox_nucleus_detector_"
+                    ) and item.name.endswith(".zarr"):
+                        shutil.rmtree(item)
+                        logger.info(
+                            "Temporary zarr directory %s has been removed.", item
+                        )
+
+        return output