refactor code and improve typing

Jiaqi-Lv · Jiaqi-Lv · commit d63a7cc05660 · 2025-12-10T16:12:14.000Z
diff --git a/tests/engines/test_nucleus_detection_engine.py b/tests/engines/test_nucleus_detection_engine.py
@@ -148,6 +148,7 @@ def test_nucleus_detector_patch(
     wsi_reader = WSIReader.open(mini_wsi_svs)
     patch_1 = wsi_reader.read_rect((0, 0), (252, 252), resolution=0.5, units="mpp")
     patch_2 = wsi_reader.read_rect((252, 252), (252, 252), resolution=0.5, units="mpp")
+    patch_3 = np.zeros((252, 252, 3), dtype=np.uint8)
 
     pretrained_model = "mapde-conic"
 
@@ -159,7 +160,7 @@ def test_nucleus_detector_patch(
         device=device,
         output_type="annotationstore",
         memory_threshold=50,
-        images=[patch_1, patch_2],
+        images=[patch_1, patch_2, patch_3],
         save_dir=save_dir,
         overwrite=True,
         class_dict=None,
@@ -173,6 +174,10 @@ def test_nucleus_detector_patch(
     assert len(store_2.values()) == 52
     store_2.close()
 
+    store_3 = SQLiteStore.open(save_dir / "2.db")
+    assert len(store_3.values()) == 0
+    store_3.close()
+
     imwrite(save_dir / "patch_0.png", patch_1)
     imwrite(save_dir / "patch_1.png", patch_2)
     _ = nucleus_detector.run(
diff --git a/tiatoolbox/data/pretrained_model.yaml b/tiatoolbox/data/pretrained_model.yaml
@@ -815,7 +815,7 @@ mapde-crchisto:
             threshold_abs: 250
             num_classes: 1
             postproc_tile_shape: [ 2048, 2048 ]
-            output_class_dict: {
+            class_dict: {
                 0: "nucleus"
             }
     ioconfig:
@@ -840,7 +840,7 @@ mapde-conic:
             threshold_abs: 205
             num_classes: 1
             postproc_tile_shape: [ 2048, 2048 ]
-            output_class_dict: {
+            class_dict: {
                 0: "nucleus"
             }
     ioconfig:
@@ -866,7 +866,7 @@ sccnn-crchisto:
             threshold_abs: 0.20
             patch_output_shape: [ 13, 13 ]
             postproc_tile_shape: [ 2048, 2048 ]
-            output_class_dict: {
+            class_dict: {
                 0: "nucleus"
             }
     ioconfig:
@@ -892,7 +892,7 @@ sccnn-conic:
             threshold_abs: 0.05
             patch_output_shape: [ 13, 13 ]
             postproc_tile_shape: [ 2048, 2048 ]
-            output_class_dict: {
+            class_dict: {
                 0: "nucleus"
             }
     ioconfig:
diff --git a/tiatoolbox/models/architecture/mapde.py b/tiatoolbox/models/architecture/mapde.py
@@ -11,9 +11,9 @@
 import numpy as np
 import torch
 import torch.nn.functional as F  # noqa: N812
-from skimage.feature import peak_local_max
 
 from tiatoolbox.models.architecture.micronet import MicroNet
+from tiatoolbox.models.architecture.utils import peak_detection_da_map_overlap
 
 
 class MapDe(MicroNet):
@@ -79,15 +79,15 @@ def __init__(
         threshold_abs: float = 250,
         num_classes: int = 1,
         postproc_tile_shape: tuple[int, int] = (2048, 2048),
-        output_class_dict: dict[int, str] | None = None,
+        class_dict: dict[int, str] | None = None,
     ) -> None:
         """Initialize :class:`MapDe`."""
         super().__init__(
             num_output_channels=num_classes * 2,
             num_input_channels=num_input_channels,
             out_activation="relu",
         )
-        self.output_class_dict = output_class_dict
+        self.output_class_dict = class_dict
         self.postproc_tile_shape = postproc_tile_shape
 
         dist_filter = np.array(
@@ -249,11 +249,6 @@ def postproc(
         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 the input block
 
         Args:
@@ -268,40 +263,14 @@ def postproc(
         Returns:
             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 ---
-        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
-
-        out = np.zeros((block_height, block_width, block_channels), dtype=np.float32)
-
-        for ch in range(block_channels):
-            img = np.asarray(block[..., ch])  # NumPy 2D view
-
-            coords = peak_local_max(
-                img,
-                min_distance=self.min_distance,
-                threshold_abs=self.threshold_abs,
-                exclude_border=False,
-            )
-
-            for r, c in coords:
-                if (rmin <= r < rmax) and (cmin <= c < cmax):
-                    out[r, c, ch] = 1.0
-
-        return out
+        return peak_detection_da_map_overlap(
+            block,
+            min_distance=self.min_distance,
+            threshold_abs=self.threshold_abs,
+            block_info=block_info,
+            depth_h=depth_h,
+            depth_w=depth_w,
+        )
 
     @staticmethod
     def infer_batch(
diff --git a/tiatoolbox/models/architecture/sccnn.py b/tiatoolbox/models/architecture/sccnn.py
@@ -10,12 +10,15 @@
 from __future__ import annotations
 
 from collections import OrderedDict
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
 
-import numpy as np
 import torch
-from skimage.feature import peak_local_max
 from torch import nn
 
+from tiatoolbox.models.architecture.utils import peak_detection_da_map_overlap
 from tiatoolbox.models.models_abc import ModelABC
 
 
@@ -92,7 +95,7 @@ def __init__(
         min_distance: int = 6,
         threshold_abs: float = 0.20,
         postproc_tile_shape: tuple[int, int] = (2048, 2048),
-        output_class_dict: dict[int, str] | None = None,
+        class_dict: dict[int, str] | None = None,
     ) -> None:
         """Initialize :class:`SCCNN`."""
         super().__init__()
@@ -102,7 +105,7 @@ def __init__(
         self.out_height = out_height
         self.out_width = out_width
         self.postproc_tile_shape = postproc_tile_shape
-        self.output_class_dict = output_class_dict
+        self.output_class_dict = class_dict
 
         # Create mesh grid and convert to 3D vector
         x, y = torch.meshgrid(
@@ -341,11 +344,6 @@ def postproc(
         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 the input block
 
         Args:
@@ -360,40 +358,14 @@ def postproc(
         Returns:
             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 ---
-        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
-
-        out = np.zeros((block_height, block_width, block_channels), dtype=np.float32)
-
-        for ch in range(block_channels):
-            img = np.asarray(block[..., ch])  # NumPy 2D view
-
-            coords = peak_local_max(
-                img,
-                min_distance=self.min_distance,
-                threshold_abs=self.threshold_abs,
-                exclude_border=False,
-            )
-
-            for r, c in coords:
-                if (rmin <= r < rmax) and (cmin <= c < cmax):
-                    out[r, c, ch] = 1.0
-
-        return out
+        return peak_detection_da_map_overlap(
+            block,
+            min_distance=self.min_distance,
+            threshold_abs=self.threshold_abs,
+            block_info=block_info,
+            depth_h=depth_h,
+            depth_w=depth_w,
+        )
 
     @staticmethod
     def infer_batch(
diff --git a/tiatoolbox/models/architecture/utils.py b/tiatoolbox/models/architecture/utils.py
@@ -7,6 +7,7 @@
 
 import numpy as np
 import torch
+from skimage.feature import peak_local_max
 from torch import nn
 
 from tiatoolbox import logger
@@ -251,3 +252,74 @@ def argmax_last_axis(image: np.ndarray) -> np.ndarray:
 
     """
     return image.argmax(axis=-1)
+
+
+def peak_detection_da_map_overlap(
+    block: np.ndarray,
+    min_distance: int,
+    threshold_abs: float | None = None,
+    threshold_rel: float | None = None,
+    block_info: dict | None = None,
+    depth_h: int = 0,
+    depth_w: int = 0,
+) -> np.ndarray:
+    """Post-processing function for peak detection.
+
+    Builds a processed mask per input channel. Runs peak_local_max then
+    writes 1.0 at peak pixels.
+
+    Can be called from Dask.da.map_overlap on a padded NumPy block
+    (h_pad, w_pad, C) to process large prediction maps in chunks with overlap.
+    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 the input block
+
+    Args:
+        block: NumPy array (H, W, C).
+        min_distance: Minimum number of pixels separating peaks.
+        threshold_abs: Minimum intensity of peaks. By default, None.
+        threshold_rel: Minimum relative intensity of peaks. By default, None.
+        block_info: Dask block info dict.
+            Only used when called from dask.array.map_overlap.
+        depth_h: Halo size in pixels for height (rows).
+            Only used when called from dask.array.map_overlap.
+        depth_w: Halo size in pixels for width (cols).
+            Only used when it's called from dask.array.map_overlap.
+
+    Returns:
+        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 ---
+    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
+
+    out = np.zeros((block_height, block_width, block_channels), dtype=np.float32)
+
+    for ch in range(block_channels):
+        img = np.asarray(block[..., ch])  # NumPy 2D view
+
+        coords = peak_local_max(
+            img,
+            min_distance=min_distance,
+            threshold_abs=threshold_abs,
+            threshold_rel=threshold_rel,
+            exclude_border=False,
+        )
+
+        for r, c in coords:
+            if (rmin <= r < rmax) and (cmin <= c < cmax):
+                out[r, c, ch] = 1.0
+
+    return out
diff --git a/tiatoolbox/models/engine/nucleus_detector.py b/tiatoolbox/models/engine/nucleus_detector.py
