✅ Add test for dask length threshold

Author: shaneahmed

tests/engines/test_semantic_segmentor.py

@@ -246,6 +246,32 @@ def test_wsi_segmentor_zarr(
     assert "count" not in output_
     assert "Current Memory usage:" in caplog.text
 
+    segmentor = SemanticSegmentor(
+        model="fcn-tissue_mask",
+        batch_size=64,
+        verbose=False,
+        num_workers=1,
+    )
+    # Return Probabilities is False
+    output = segmentor.run(
+        images=[sample_svs],
+        return_probabilities=True,
+        return_labels=False,
+        device=device,
+        patch_mode=False,
+        save_dir=tmp_path / "task_length_cache",
+        batch_size=2,
+        output_type="zarr",
+        da_length_threshold=1,
+    )
+
+    output_ = zarr.open(output[sample_svs], mode="r")
+    assert 0.17 < np.mean(output_["predictions"][:]) < 0.19
+    assert "probabilities" in output_
+    assert "canvas" not in output_
+    assert "count" not in output_
+    assert "Canvas task graph length:" in caplog.text
+
     # Return Probabilities is True
     # Using small image for faster run
     segmentor = SemanticSegmentor(

tiatoolbox/models/engine/engine_abc.py

@@ -58,6 +58,7 @@
     dict_to_store_patch_predictions,
     get_tqdm,
 )
+from tiatoolbox.wsicore.wsireader import is_zarr
 
 from .io_config import ModelIOConfigABC
 
@@ -644,14 +645,16 @@ def save_predictions(
         keys_to_compute = [k for k in processed_predictions if k not in self.drop_keys]
 
         if output_type.lower() == "zarr":
+            if is_zarr(save_path):
+                zarr_group = zarr.open(save_path, mode="r")
+                keys_to_compute = [k for k in keys_to_compute if k not in zarr_group]
             write_tasks = []
             for key in keys_to_compute:
                 dask_array = processed_predictions[key]
                 task = dask_array.to_zarr(
                     url=save_path,
                     component=key,
                     compute=False,
-                    overwrite=True,
                 )
                 write_tasks.append(task)
             msg = f"Saving output to {save_path}."