✅ Test patch mode with dict output

Author: shaneahmed

requirements/requirements.txt

@@ -4,7 +4,8 @@ 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
+dask[array]>=2025.10.0
+dask[dataframe]>=2025.10.0
 defusedxml>=0.7.1
 filelock>=3.9.0
 flask>=2.2.2

tests/engines/test_nucleus_instance_segmentor.py

@@ -1,20 +1,19 @@
 """Test tiatoolbox.models.engine.nucleus_instance_segmentor."""
 
-import shutil
 from collections.abc import Callable
 from pathlib import Path
-from typing import Literal, Final
+from typing import Final
 
-import torch
 import numpy as np
+import torch
 
-from tiatoolbox.models import IOSegmentorConfig, NucleusInstanceSegmentor
+from tiatoolbox.models import NucleusInstanceSegmentor
 from tiatoolbox.wsicore import WSIReader
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 
-def test_functionality_patch_mode(remote_sample: Callable, track_tmp_path: Path) -> None:
+def test_functionality_patch_mode(remote_sample: Callable) -> None:
     """Patch mode functionality test for nuclei instance segmentor."""
     mini_wsi_svs = Path(remote_sample("wsi4_1k_1k_svs"))
     mini_wsi = WSIReader.open(mini_wsi_svs)
@@ -34,10 +33,7 @@ def test_functionality_patch_mode(remote_sample: Callable, track_tmp_path: Path)
         units=units,
     )
 
-    patches = np.stack(
-        arrays=[patch1, patch2],
-        axis=0
-    )
+    patches = np.stack(arrays=[patch1, patch2], axis=0)
 
     inst_segmentor = NucleusInstanceSegmentor(
         batch_size=1,
@@ -48,29 +44,10 @@ def test_functionality_patch_mode(remote_sample: Callable, track_tmp_path: Path)
         images=patches,
         patch_mode=True,
         device=device,
-        save_dir=track_tmp_path / "hovernet_fast-pannuke",
         output_type="dict",
     )
 
-    assert output
-
-
-def test_functionality_wsi(remote_sample: Callable, track_tmp_path: Path) -> None:
-    """Local functionality test for nuclei instance segmentor."""
-    root_save_dir = Path(track_tmp_path)
-    save_dir = Path(f"{track_tmp_path}/output")
-    mini_wsi_svs = Path(remote_sample("wsi4_1k_1k_svs"))
-
-    # * generate full output w/o parallel post-processing worker first
-    shutil.rmtree(save_dir, ignore_errors=True)
-    inst_segmentor = NucleusInstanceSegmentor(
-        batch_size=8,
-        num_postproc_workers=0,
-        pretrained_model="hovernet_fast-pannuke",
-    )
-    output = inst_segmentor.run(
-        [mini_wsi_svs],
-        patch_mode=False,
-        device=device,
-        save_dir=save_dir,
-    )
+    assert np.max(output["predictions"][0][:]) == 41
+    assert np.max(output["predictions"][1][:]) == 17
+    assert len(output["inst_dict"][0].columns) == 41
+    assert len(output["inst_dict"][1].columns) == 17

tiatoolbox/models/engine/nucleus_instance_segmentor.py

@@ -4,40 +4,41 @@
 
 import uuid
 from collections import deque
+from pathlib import Path
 from typing import TYPE_CHECKING
 
-import dask
+import dask.array as da
+import dask.dataframe as dd
+
 # replace with the sql database once the PR in place
 import joblib
 import numpy as np
+import pandas as pd
 import torch
 import tqdm
-import dask.array as da
 from shapely.geometry import box as shapely_box
 from shapely.strtree import STRtree
-from torch.utils.data import DataLoader
 from typing_extensions import Unpack
 
+from tiatoolbox import DuplicateFilter, logger
 from tiatoolbox.models.engine.semantic_segmentor import (
     SemanticSegmentor,
     SemanticSegmentorRunParams,
 )
 from tiatoolbox.tools.patchextraction import PatchExtractor
-from tiatoolbox.models.models_abc import ModelABC
 from tiatoolbox.utils.misc import get_tqdm
-from .engine_abc import EngineABCRunParams
-from tiatoolbox import DuplicateFilter, logger
-from pathlib import Path
-
 
 if TYPE_CHECKING:  # pragma: no cover
     import os
     from collections.abc import Callable
 
+    from torch.utils.data import DataLoader
 
     from tiatoolbox.annotation import AnnotationStore
+    from tiatoolbox.models.models_abc import ModelABC
     from tiatoolbox.wsicore import WSIReader
 
+    from .engine_abc import EngineABCRunParams
     from .io_config import IOInstanceSegmentorConfig, IOSegmentorConfig
 
 
@@ -490,7 +491,9 @@ def infer_patches(
                 labels.append(da.from_array(np.array(batch_data["label"])))
 
         for i in range(num_expected_output):
-            raw_predictions["probabilities"][i] = da.concatenate(probabilities[i], axis=0)
+            raw_predictions["probabilities"][i] = da.concatenate(
+                probabilities[i], axis=0
+            )
 
         if return_coordinates:
             raw_predictions["coordinates"] = da.concatenate(coordinates, axis=0)
@@ -548,8 +551,8 @@ def _run_patch_mode(
             return_coordinates=output_type == "annotationstore",
         )
 
-        raw_predictions["predictions"] = self.post_process_patches(
-            raw_predictions=raw_predictions["probabilities"],
+        raw_predictions = self.post_process_patches(
+            raw_predictions=raw_predictions,
             prediction_shape=None,
             prediction_dtype=None,
             **kwargs,
@@ -570,11 +573,11 @@ def _run_patch_mode(
 
     def post_process_patches(  # skipcq: PYL-R0201
         self: NucleusInstanceSegmentor,
-        raw_predictions: da.Array,
+        raw_predictions: dict,
         prediction_shape: tuple[int, ...],  # noqa: ARG002
         prediction_dtype: type,  # noqa: ARG002
         **kwargs: Unpack[EngineABCRunParams],  # noqa: ARG002
-    ) -> dask.array.Array:
+    ) -> dict:
         """Post-process raw patch predictions from inference.
 
         This method applies a post-processing function (e.g., smoothing, filtering)
@@ -596,9 +599,44 @@ def post_process_patches(  # skipcq: PYL-R0201
                 Post-processed predictions as a Dask array.
 
         """
-        raw_predictions = self.model.postproc_func(raw_predictions)
+        probabilities = raw_predictions["probabilities"]
+        predictions = [[] for _ in range(probabilities[0].shape[0])]
+        inst_dict = [[] for _ in range(probabilities[0].shape[0])]
+        for idx in range(probabilities[0].shape[0]):
+            predictions[idx], inst_dict[idx] = self.model.postproc_func(
+                [probabilities[0][idx], probabilities[1][idx], probabilities[2][idx]]
+            )
+            inst_dict[idx] = dd.from_pandas(pd.DataFrame(inst_dict[idx]))
+
+        raw_predictions["predictions"] = da.stack(predictions, axis=0)
+        raw_predictions["inst_dict"] = inst_dict
+
         return raw_predictions
 
+    def save_predictions(
+        self: SemanticSegmentor,
+        processed_predictions: dict,
+        output_type: str,
+        save_path: Path | None = None,
+        **kwargs: Unpack[SemanticSegmentorRunParams],
+    ) -> dict | AnnotationStore | Path:
+        """Save semantic segmentation predictions to disk or return them in memory."""
+        # Conversion to annotationstore uses a different function for SemanticSegmentor
+        inst_dict: list[dd.DataFrame] | None = processed_predictions.pop(
+            "inst_dict", None
+        )
+        out = super().save_predictions(
+            processed_predictions, output_type, save_path=save_path, **kwargs
+        )
+
+        if isinstance(out, dict):
+            out["inst_dict"] = [[] for _ in range(len(inst_dict))]
+            for idx in range(len(inst_dict)):
+                out["inst_dict"][idx] = inst_dict[idx].compute()
+            return out
+
+        return out
+
     @staticmethod
     def _get_tile_info(
         image_shape: list[int] | np.ndarray,