✅ Test patch mode with dict and zarr output

Author: shaneahmed

tests/engines/test_nucleus_instance_segmentor.py

@@ -6,14 +6,17 @@
 
 import numpy as np
 import torch
+import zarr
 
 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) -> None:
+def test_functionality_patch_mode(
+    remote_sample: Callable, track_tmp_path: Path
+) -> 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)
@@ -33,7 +36,10 @@ def test_functionality_patch_mode(remote_sample: Callable) -> None:
         units=units,
     )
 
-    patches = np.stack(arrays=[patch1, patch2], axis=0)
+    # Test dummy input, should result in no output segmentation
+    patch3 = np.zeros_like(patch1)
+
+    patches = np.stack(arrays=[patch1, patch2, patch3], axis=0)
 
     inst_segmentor = NucleusInstanceSegmentor(
         batch_size=1,
@@ -49,5 +55,85 @@ def test_functionality_patch_mode(remote_sample: Callable) -> None:
 
     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
+    assert np.max(output["predictions"][2][:]) == 0
+
+    assert len(output["box"][0]) == 41
+    assert len(output["box"][1]) == 17
+    assert len(output["box"][2]) == 0
+
+    assert len(output["centroid"][0]) == 41
+    assert len(output["centroid"][1]) == 17
+    assert len(output["centroid"][2]) == 0
+
+    assert len(output["contour"][0]) == 41
+    assert len(output["contour"][1]) == 17
+    assert len(output["contour"][2]) == 0
+
+    assert len(output["prob"][0]) == 41
+    assert len(output["prob"][1]) == 17
+    assert len(output["prob"][2]) == 0
+
+    assert len(output["type"][0]) == 41
+    assert len(output["type"][1]) == 17
+    assert len(output["type"][2]) == 0
+
+    output_ = output
+
+    output = inst_segmentor.run(
+        images=patches,
+        patch_mode=True,
+        device=device,
+        output_type="zarr",
+        save_dir=track_tmp_path / "patch_output_zarr",
+    )
+
+    output = zarr.open(output, mode="r")
+
+    assert np.max(output["predictions"][0][:]) == 41
+    assert np.max(output["predictions"][1][:]) == 17
+
+    assert all(
+        np.array_equal(a, b)
+        for a, b in zip(output["box"][0], output_["box"][0], strict=False)
+    )
+    assert all(
+        np.array_equal(a, b)
+        for a, b in zip(output["box"][1], output_["box"][1], strict=False)
+    )
+    assert len(output["box"][2]) == 0
+
+    assert all(
+        np.array_equal(a, b)
+        for a, b in zip(output["centroid"][0], output_["centroid"][0], strict=False)
+    )
+    assert all(
+        np.array_equal(a, b)
+        for a, b in zip(output["centroid"][1], output_["centroid"][1], strict=False)
+    )
+
+    assert all(
+        np.array_equal(a, b)
+        for a, b in zip(output["contour"][0], output_["contour"][0], strict=False)
+    )
+    assert all(
+        np.array_equal(a, b)
+        for a, b in zip(output["contour"][1], output_["contour"][1], strict=False)
+    )
+
+    assert all(
+        np.array_equal(a, b)
+        for a, b in zip(output["prob"][0], output_["prob"][0], strict=False)
+    )
+    assert all(
+        np.array_equal(a, b)
+        for a, b in zip(output["prob"][1], output_["prob"][1], strict=False)
+    )
+
+    assert all(
+        np.array_equal(a, b)
+        for a, b in zip(output["type"][0], output_["type"][0], strict=False)
+    )
+    assert all(
+        np.array_equal(a, b)
+        for a, b in zip(output["type"][1], output_["type"][1], strict=False)
+    )

tiatoolbox/models/architecture/hovernet.py

@@ -4,9 +4,13 @@
 
 import math
 from collections import OrderedDict
-import dask
+
 import cv2
+import dask
+import dask.array as da
+import dask.dataframe as dd
 import numpy as np
+import pandas as pd
 import torch
 import torch.nn.functional as F  # noqa: N812
 from scipy import ndimage
@@ -22,6 +26,8 @@
 from tiatoolbox.models.models_abc import ModelABC
 from tiatoolbox.utils.misc import get_bounding_box
 
+dask.config.set({"dataframe.convert-string": False})
+
 
 class TFSamepaddingLayer(nn.Module):
     """To align with tensorflow `same` padding.
@@ -782,7 +788,28 @@ def postproc(raw_maps: list[np.ndarray]) -> tuple[np.ndarray, dict]:
         pred_inst = HoVerNet._proc_np_hv(np_map, hv_map)
         nuc_inst_info_dict = HoVerNet.get_instance_info(pred_inst, pred_type)
 
-        return pred_inst, nuc_inst_info_dict
+        if not nuc_inst_info_dict:
+            nuc_inst_info_dict = {  # inst_id should start at 1
+                "box": da.empty(shape=0),
+                "centroid": da.empty(shape=0),
+                "contour": da.empty(shape=0),
+                "prob": da.empty(shape=0),
+                "type": da.empty(shape=0),
+            }
+            return pred_inst, nuc_inst_info_dict
+
+        # dask dataframe does not support transpose
+        nuc_inst_info_df = pd.DataFrame(nuc_inst_info_dict).transpose()
+
+        # create dask dataframe
+        nuc_inst_info_dd = dd.from_pandas(nuc_inst_info_df)
+
+        # reinitialize nuc_inst_info_dict
+        nuc_inst_info_dict_ = {}
+        for key in nuc_inst_info_df.columns:
+            nuc_inst_info_dict_[key] = nuc_inst_info_dd[key].to_dask_array(lengths=True)
+
+        return pred_inst, nuc_inst_info_dict_
 
     @staticmethod
     def infer_batch(  # skipcq: PYL-W0221

tiatoolbox/models/engine/engine_abc.py

@@ -46,6 +46,7 @@
 import zarr
 from dask import compute
 from dask.diagnostics import ProgressBar
+from numcodecs import Pickle
 from torch import nn
 from typing_extensions import Unpack
 
@@ -71,6 +72,8 @@
     from tiatoolbox.models.models_abc import ModelABC
     from tiatoolbox.type_hints import IntPair, Resolution, Units
 
+dask.config.set({"dataframe.convert-string": False})
+
 
 class EngineABCRunParams(TypedDict, total=False):
     """Parameters for configuring the :func:`EngineABC.run()` method.
@@ -645,13 +648,29 @@ def save_predictions(
                 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].rechunk("auto")
-                task = dask_array.to_zarr(
-                    url=save_path,
-                    component=key,
-                    compute=False,
-                )
-                write_tasks.append(task)
+                dask_output = processed_predictions[key]
+                if isinstance(dask_output, da.Array):
+                    dask_output = dask_output.rechunk("auto")
+                    task = dask_output.to_zarr(
+                        url=save_path, component=key, compute=False, object_codec=None
+                    )
+                    write_tasks.append(task)
+
+                if isinstance(dask_output, list) and all(
+                    isinstance(dask_array, da.Array) for dask_array in dask_output
+                ):
+                    for i, dask_array in enumerate(dask_output):
+                        object_codec = (
+                            Pickle() if dask_array.dtype == "object" else None
+                        )
+                        task = dask_array.to_zarr(
+                            url=save_path,
+                            component=f"{key}/{i}",
+                            compute=False,
+                            object_codec=object_codec,
+                        )
+                        write_tasks.append(task)
+
             msg = f"Saving output to {save_path}."
             logger.info(msg=msg)
             with ProgressBar():

tiatoolbox/models/engine/nucleus_instance_segmentor.py

@@ -8,12 +8,10 @@
 from typing import TYPE_CHECKING
 
 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
 from shapely.geometry import box as shapely_box
@@ -601,15 +599,15 @@ def post_process_patches(  # skipcq: PYL-R0201
         """
         probabilities = raw_predictions["probabilities"]
         predictions = [[] for _ in range(probabilities[0].shape[0])]
-        inst_dict = [[] 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
+        for key in inst_dict[0]:
+            raw_predictions[key] = [d[key] for d in inst_dict]
 
         return raw_predictions
 
@@ -621,22 +619,10 @@ def save_predictions(
         **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(
+        return 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,