update patch mode processing

Author: Jiaqi-Lv

test.py

@@ -1,26 +1,46 @@
-import pathlib
+from pathlib import Path
 
 from tiatoolbox.models.engine.nucleus_detector import (
     NucleusDetector,
 )
 from tiatoolbox.utils import env_detection as toolbox_env
+from tiatoolbox.wsicore.wsireader import WSIReader
 
-ON_GPU = not toolbox_env.running_on_ci() and toolbox_env.has_gpu()
+from tiatoolbox.utils.misc import imwrite
 
+ON_GPU = not toolbox_env.running_on_ci() and toolbox_env.has_gpu()
 
 if __name__ == "__main__":
-    # model_name = "sccnn-crchisto"
-    model_name = "mapde-conic"
+    model_name = "sccnn-crchisto"
+    # model_name = "mapde-conic"
+
+
+    # test_image_path = "/media/u1910100/data/slides/CMU-1-Small-Region.svs"
+    # reader = WSIReader.open(test_image_path)
+
+    # patch_1 = reader.read_region((1500, 1500), level=0, size=(31, 31))
+
+    # imwrite(Path("/media/u1910100/data/slides/patch_1.png"), patch_1)
+
+    # patch_2 = reader.read_region((1000, 1000), level=0, size=(31, 31))
+    # imwrite(Path("/media/u1910100/data/slides/patch_2.png"), patch_2)
+
+    # patches = [
+    #     Path("/media/u1910100/data/slides/patch_1.png"),
+    #     Path("/media/u1910100/data/slides/patch_2.png"),
+    # ]
+
 
     detector = NucleusDetector(model=model_name, batch_size=16, num_workers=8)
     detector.run(
-        images=[pathlib.Path("/media/u1910100/data/slides/CMU-1-Small-Region.svs")],
+        images=[Path("/media/u1910100/data/slides/wsi1_2k_2k.svs")],
+        # images=patches,
         patch_mode=False,
         device="cuda",
-        save_dir=pathlib.Path("/media/u1910100/data/overlays/test"),
+        save_dir=Path("/media/u1910100/data/overlays/test"),
         overwrite=True,
         output_type="annotationstore",
         class_dict={0: "nucleus"},
         auto_get_mask=True,
-        memory_threshold=80,
+        memory_threshold=70,
     )

tiatoolbox/data/pretrained_model.yaml

@@ -867,7 +867,6 @@ sccnn-crchisto:
                 - { "units": "mpp", "resolution": 0.25 }
             output_resolutions:
                 - { "units": "mpp", "resolution": 0.25 }
-            tile_shape: [ 2048, 2048 ]
             patch_input_shape: [ 31, 31 ]
             patch_output_shape: [ 13, 13 ]
             stride_shape: [ 8, 8 ]
@@ -891,7 +890,6 @@ sccnn-conic:
                 - { "units": "mpp", "resolution": 0.25 }
             output_resolutions:
                 - { "units": "mpp", "resolution": 0.25 }
-            tile_shape: [ 2048, 2048 ]
             patch_input_shape: [ 31, 31 ]
             patch_output_shape: [ 13, 13 ]
             stride_shape: [ 8, 8 ]

tiatoolbox/models/architecture/mapde.py

@@ -239,34 +239,44 @@ def forward(self: MapDe, input_tensor: torch.Tensor) -> torch.Tensor:
     def postproc(
         self: MapDe,
         block: np.ndarray,
-        block_info: dict,
-        depth_h: int,
-        depth_w: int,
+        block_info: dict | None = None,
+        depth_h: int = 0,
+        depth_w: int = 0,
     ) -> np.ndarray:
-        """Runs inside Dask.da.map_overlap on a padded NumPy block: (h_pad, w_pad, C).
+        """ MapDe post-processing function.
 
-        Builds a processed mask per channel, runs peak_local_max then
-        writes 1.0 at centroid pixels.
-        Keeps only centroids whose (row,col) lie in the interior window:
+        Builds a processed mask per input channel, runs peak_local_max then
+        writes 1.0 at peak pixels.
+
+        Can be called inside Dask.da.map_overlap on a padded NumPy block: (h_pad, w_pad, C)
+        to process large prediction maps in chunks. Keeps only centroids whose (row,col) 
+        lie in the interior window:
             rows [depth_h : depth_h + core_h), cols [depth_w : depth_w + core_w)
-        Returns same spatial shape as input block: (h_pad, w_pad, C), float32.
+
+        Returns same spatial shape as the input block
 
         Args:
-            block: NumPy array (H, W, C) with padded block data.
-            block_info: Dask block info dict.
-            depth_h: Halo size in pixels for height (rows).
-            depth_w: Halo size in pixels for width (cols).
+            block: NumPy array (H, W, C).
+            block_info: Dask block info dict. Only used when called inside dask.array.map_overlap.
+            depth_h: Halo size in pixels for height (rows). 
+                Only used when it's called inside dask.array.map_overlap.
+            depth_w: Halo size in pixels for width (cols). 
+                Only used when it's called inside dask.array.map_overlap.
 
         Returns:
-            out: NumPy array (H, W, C) with 1 at centroids, 0 elsewhere.
+            out: NumPy array (H, W, C) with 1.0 at peaks, 0 elsewhere.
         """
         block_height, block_width, block_channels = block.shape
 
         # --- derive core (pre-overlap) size for THIS block ---
-        info = block_info[0]
-        locs = info["array-location"]  # a list of (start, stop) coordinates per axis
-        core_h = int(locs[0][1] - locs[0][0])  # r1 - r0
-        core_w = int(locs[1][1] - locs[1][0])
+        if block_info is None:
+            core_h = block_height - 2 * depth_h
+            core_w = block_width - 2 * depth_w
+        else:
+            info = block_info[0]
+            locs = info["array-location"]  # a list of (start, stop) coordinates per axis
+            core_h = int(locs[0][1] - locs[0][0])  # r1 - r0
+            core_w = int(locs[1][1] - locs[1][0])
 
         rmin, rmax = depth_h, depth_h + core_h
         cmin, cmax = depth_w, depth_w + core_w

tiatoolbox/models/architecture/sccnn.py

@@ -329,34 +329,44 @@ def spatially_constrained_layer1(
     def postproc(
         self: SCCNN,
         block: np.ndarray,
-        block_info: dict,
-        depth_h: int,
-        depth_w: int,
+        block_info: dict | None = None,
+        depth_h: int = 0,
+        depth_w: int = 0,
     ) -> np.ndarray:
-        """Runs inside Dask.da.map_overlap on a padded NumPy block: (h_pad, w_pad, C).
+        """ SCCNN post-processing function.
 
-        Builds a processed mask per channel, runs peak_local_max then
-        writes 1.0 at centroid pixels.
-        Keeps only centroids whose (row,col) lie in the interior window:
+        Builds a processed mask per input channel, runs peak_local_max then
+        writes 1.0 at peak pixels.
+
+        Can be called inside Dask.da.map_overlap on a padded NumPy block: (h_pad, w_pad, C)
+        to process large prediction maps in chunks. Keeps only centroids whose (row,col) 
+        lie in the interior window:
             rows [depth_h : depth_h + core_h), cols [depth_w : depth_w + core_w)
-        Returns same spatial shape as input block: (h_pad, w_pad, C), float32.
+
+        Returns same spatial shape as the input block
 
         Args:
-            block: NumPy array (H, W, C) with padded block data.
-            block_info: Dask block info dict.
-            depth_h: Halo size in pixels for height (rows).
-            depth_w: Halo size in pixels for width (cols).
+            block: NumPy array (H, W, C).
+            block_info: Dask block info dict. Only used when called inside dask.array.map_overlap.
+            depth_h: Halo size in pixels for height (rows). 
+                Only used when it's called inside dask.array.map_overlap.
+            depth_w: Halo size in pixels for width (cols). 
+                Only used when it's called inside dask.array.map_overlap.
 
         Returns:
-            out: NumPy array (H, W, C) with 1 at centroids, 0 elsewhere.
+            out: NumPy array (H, W, C) with 1.0 at peaks, 0 elsewhere.
         """
         block_height, block_width, block_channels = block.shape
 
         # --- derive core (pre-overlap) size for THIS block ---
-        info = block_info[0]
-        locs = info["array-location"]  # a list of (start, stop) coordinates per axis
-        core_h = int(locs[0][1] - locs[0][0])  # r1 - r0
-        core_w = int(locs[1][1] - locs[1][0])
+        if block_info is None:
+            core_h = block_height - 2 * depth_h
+            core_w = block_width - 2 * depth_w
+        else:
+            info = block_info[0]
+            locs = info["array-location"]  # a list of (start, stop) coordinates per axis
+            core_h = int(locs[0][1] - locs[0][0])  # r1 - r0
+            core_w = int(locs[1][1] - locs[1][0])
 
         rmin, rmax = depth_h, depth_h + core_h
         cmin, cmax = depth_w, depth_w + core_w

tiatoolbox/models/engine/nucleus_detector.py

@@ -96,11 +96,11 @@ def __init__(
 
     def post_process_patches(
         self: NucleusDetector,
-        raw_predictions: da.Array,
+        raw_predictions: list[da.Array],
         prediction_shape: tuple[int, ...],
         prediction_dtype: type,
         **kwargs: Unpack[SemanticSegmentorRunParams],
-    ) -> list[pd.DataFrame]:
+    ) -> list[np.ndarray]:
         """Define how to post-process patch predictions.
 
         Args:
@@ -117,7 +117,7 @@ def post_process_patches(
         _ = prediction_dtype
 
         batch_predictions = []
-        for i in range(raw_predictions.shape[0]):
+        for i in range(len(raw_predictions)):
             batch_predictions.append(self.model.postproc_func(raw_predictions[i]))
         return batch_predictions
 
@@ -233,13 +233,15 @@ def save_predictions(
         if self.patch_mode:
             save_paths = []
             for i, predictions in enumerate(processed_predictions["predictions"]):
+                predictions_da = da.from_array(predictions, chunks=predictions.shape)
+
                 if isinstance(self.images[i], Path):
                     output_path = save_path.parent / (self.images[i].stem + ".db")
                 else:
                     output_path = save_path.parent / (str(i) + ".db")
 
                 out_file = self.write_centroids_to_store(
-                    predictions,
+                    predictions_da,
                     scale_factor=scale_factor,
                     class_dict=class_dict,
                     save_path=output_path,