refactor postproc to accept 'raw_prediictions'

Author: Jiaqi Lv

tests/engines/test_engine_abc.py

@@ -8,6 +8,7 @@
 from pathlib import Path
 from typing import NoReturn
 
+import dask.array as da
 import numpy as np
 import pytest
 import torch
@@ -353,11 +354,13 @@ def test_engine_run() -> NoReturn:
     assert "probabilities" in out
     assert "labels" in out
 
+    raw_output = np.zeros((3, 3, 3))
     pred = eng.post_process_wsi(
-        raw_predictions=Path("/path/to/raw_predictions.npy"),
+        raw_predictions={"probabilities": da.from_array(raw_output)},
         save_path=Path("/path/to/save_predictions.zarr"),
     )
-    assert str(pred) == "/path/to/raw_predictions.npy"
+    pred = np.array(pred)
+    np.testing.assert_array_equal(pred, raw_output)
 
 
 def test_engine_run_with_verbose() -> NoReturn:

tiatoolbox/models/engine/engine_abc.py

@@ -625,7 +625,7 @@ def post_process_patches(  # skipcq: PYL-R0201
                 Post-processed predictions as a Dask array.
 
         """
-        return raw_predictions
+        return raw_predictions["probabilities"]
 
     def save_predictions(
         self: EngineABC,
@@ -933,7 +933,7 @@ def post_process_wsi(  # skipcq: PYL-R0201
                 Post-processed predictions as a Dask array.
 
         """
-        return raw_predictions
+        return raw_predictions["probabilities"]
 
     def _load_ioconfig(self: EngineABC, ioconfig: ModelIOConfigABC) -> ModelIOConfigABC:
         """Load or validate the IO configuration for the engine.
@@ -1367,7 +1367,7 @@ def _run_patch_mode(
         )
 
         raw_predictions["predictions"] = self.post_process_patches(
-            raw_predictions=raw_predictions["probabilities"],
+            raw_predictions=raw_predictions,
             prediction_shape=raw_predictions["probabilities"].shape[:-1],
             prediction_dtype=raw_predictions["probabilities"].dtype,
             **kwargs,
@@ -1543,7 +1543,7 @@ def get_path(image: Path | WSIReader) -> Path:
             )
 
             raw_predictions["predictions"] = self.post_process_wsi(
-                raw_predictions=raw_predictions["probabilities"],
+                raw_predictions=raw_predictions,
                 save_path=save_path[get_path(image)],
                 prediction_shape=raw_predictions["probabilities"].shape[:-1],
                 prediction_dtype=raw_predictions["probabilities"].dtype,

tiatoolbox/models/engine/nucleus_detector.py

@@ -356,8 +356,8 @@ def post_process_patches(
         probs = []
 
         # Process each patch's predictions
-        for i in range(raw_predictions.shape[0]):
-            probs_prediction_patch = raw_predictions[i].compute()
+        for i in range(raw_predictions["probabilities"].shape[0]):
+            probs_prediction_patch = raw_predictions["probabilities"][i].compute()
             centroids_map_patch = self.model.postproc(
                 probs_prediction_patch,
                 min_distance=min_distance,
@@ -462,13 +462,13 @@ def post_process_wsi(
         depth = {0: depth_h, 1: depth_w, 2: 0}
 
         # Re-chunk to post-processing tile shape for more efficient processing
-        rechunked_prediction_map = raw_predictions.rechunk(
+        rechunked_probability_map = raw_predictions["probabilities"].rechunk(
             (postproc_tile_shape[0], postproc_tile_shape[1], -1)
         )
 
         centroid_maps = da.map_overlap(
             self.model.postproc,
-            rechunked_prediction_map,
+            rechunked_probability_map,
             min_distance=min_distance,
             threshold_abs=threshold_abs,
             threshold_rel=threshold_rel,

tiatoolbox/models/engine/patch_predictor.py

@@ -381,7 +381,7 @@ def post_process_patches(
         _ = kwargs.get("return_probabilities")
         _ = prediction_shape
         _ = prediction_dtype
-        raw_predictions = self.model.postproc_func(raw_predictions)
+        raw_predictions = self.model.postproc_func(raw_predictions["probabilities"])
         return cast_to_min_dtype(raw_predictions)
 
     def post_process_wsi(