✅ Check for cache_mode with zarr output.

Author: shaneahmed

tests/engines/test_semantic_segmentor.py

@@ -4,7 +4,9 @@
 
 from pathlib import Path
 
+import numpy as np
 import torch
+import zarr
 
 from tiatoolbox.models.engine.semantic_segmentor_new import SemanticSegmentor
 from tiatoolbox.utils import env_detection as toolbox_env
@@ -20,7 +22,9 @@ def test_semantic_segmentor_init() -> None:
     assert isinstance(segmentor.model, torch.nn.Module)
 
 
-def test_semantic_segmentor_patches(sample_patch1: Path, sample_patch2: Path) -> None:
+def test_semantic_segmentor_patches(
+    sample_patch1: Path, sample_patch2: Path, tmp_path: Path
+) -> None:
     """Tests SemanticSegmentor on image patches."""
     segmentor = SemanticSegmentor(
         model="fcn-tissue_mask", batch_size=32, verbose=False, device=device
@@ -38,6 +42,9 @@ def test_semantic_segmentor_patches(sample_patch1: Path, sample_patch2: Path) ->
         patch_mode=True,
     )
 
+    assert 0.24 < np.mean(output["predictions"][:]) < 0.25
+    assert 0.495 < np.mean(output["probabilities"][:]) < 0.505
+
     assert (
         tuple(segmentor._ioconfig.patch_output_shape)
         == output["probabilities"][0].shape[:-1]
@@ -46,3 +53,35 @@ def test_semantic_segmentor_patches(sample_patch1: Path, sample_patch2: Path) ->
     assert (
         tuple(segmentor._ioconfig.patch_output_shape) == output["predictions"][0].shape
     )
+
+    output = segmentor.run(
+        images=inputs,
+        return_probabilities=True,
+        return_labels=False,
+        device=device,
+        patch_mode=True,
+        cache_mode=True,
+        save_dir=tmp_path / "output0",
+    )
+
+    assert output == tmp_path / "output0" / "output.zarr"
+
+    output = zarr.open(output, mode="r")
+    assert 0.24 < np.mean(output["predictions"][:]) < 0.25
+    assert 0.495 < np.mean(output["probabilities"][:]) < 0.505
+
+    output = segmentor.run(
+        images=inputs,
+        return_probabilities=False,
+        return_labels=False,
+        device=device,
+        patch_mode=True,
+        cache_mode=True,
+        save_dir=tmp_path / "output1",
+    )
+
+    assert output == tmp_path / "output1" / "output.zarr"
+
+    output = zarr.open(output, mode="r")
+    assert 0.24 < np.mean(output["predictions"][:]) < 0.25
+    assert "probabilities" not in output.keys()  # noqa: SIM118