🆕 Add KongNet Model and Example Notebook

tests/models/test_arch_utils.py

@@ -6,9 +6,11 @@
 import torch
 
 from tiatoolbox.models.architecture.utils import (
+    AttentionModule,
     UpSample2x,
     centre_crop,
     centre_crop_to_shape,
+    nms_on_detection_maps,
     peak_detection_map_overlap,
 )
 
@@ -160,3 +162,38 @@ def test_peak_detection_map_overlap() -> None:
     assert peak_map[0, 0, 0] == 1.0
     assert peak_map[3, 3, 0] == 1.0
     assert np.sum(peak_map) == 2.0
+
+
+def test_nms_on_detection_maps() -> None:
+    """Test for NMS on detection maps."""
+    heatmap = np.zeros((7, 7, 3), dtype=np.float32)
+    nms_map = nms_on_detection_maps(heatmap, min_distance=3)
+    assert np.sum(nms_map) == 0.0  # No peaks
+
+    heatmap[0, 0, 0] = 0.9  # Peak in channel 0 (valid)
+    heatmap[0, 1, 0] = 0.6  # Peak in channel 0 (suppressed)
+    heatmap[0, 0, 1] = 0.8  # Peak in channel 1 (suppressed)
+
+    heatmap[5, 5, 2] = 0.9  # Peak in channel 2 (valid)
+    heatmap[4, 4, 1] = 0.7  # Peak in channel 1 (suppressed)
+
+    nms_map = nms_on_detection_maps(heatmap, min_distance=3)
+    assert nms_map[0, 0, 0] == 0.9
+    assert nms_map[5, 5, 2] == 0.9
+
+
+def test_attention_module() -> None:
+    """Test for Attention module."""
+    test_input = torch.zeros((1, 16, 32, 32), dtype=torch.float32)
+
+    # Default to identity
+    attention = AttentionModule(name=None, in_channels=16)
+    output = attention(test_input)
+    assert torch.sum(output - test_input) == 0
+
+    attention = AttentionModule(name="scse", in_channels=16, reduction=4)
+    output = attention(test_input)
+    assert output.shape == test_input.shape
+
+    with pytest.raises(ValueError, match=r"Attention random_name is not implemented"):
+        _ = AttentionModule(name="random_name", in_channels=16)

tiatoolbox/data/pretrained_model.yaml

@@ -964,3 +964,187 @@ grandqc_tissue_detection:
             patch_output_shape: [512, 512]
             stride_shape: [256, 256]
             save_resolution: {'units': 'mpp', 'resolution': 10.0}
+
+KongNet_CoNIC_1:
+    hf_repo_id: TIACentre/KongNet_pretrained_weights
+    architecture:
+        class: kongnet.KongNet
+        kwargs:
+            num_heads: 6
+            num_channels_per_head: [3, 3, 3, 3, 3, 3]
+            target_channels: [2, 5, 8, 11, 14, 17]
+            wide_decoder: False
+            min_distance: 5
+            threshold_abs: 0.5
+            class_dict: {
+                0: "Neutrophil",
+                1: "Epithelial",
+                2: "Lymphocyte",
+                3: "Plasma",
+                4: "Eosinophil",
+                5: "Connective",
+            }
+            postproc_tile_shape: [2048, 2048]
+    ioconfig:
+        class: io_config.IOSegmentorConfig
+        kwargs:
+            input_resolutions:
+                - {'units': 'mpp', 'resolution': 0.5}
+            output_resolutions:
+                - {'units': 'mpp', 'resolution': 0.5}
+            patch_input_shape: [256, 256]
+            patch_output_shape: [256, 256]
+            stride_shape: [248, 248]
+            save_resolution: {'units': 'baseline', 'resolution': 1.0}
+
+KongNet_PanNuke_1:
+    hf_repo_id: TIACentre/KongNet_pretrained_weights
+    architecture:
+        class: kongnet.KongNet
+        kwargs:
+            num_heads: 6
+            num_channels_per_head: [3, 3, 3, 3, 3, 3]
+            target_channels: [5, 8, 11, 14, 17]
+            wide_decoder: False
+            min_distance: 11
+            threshold_abs: 0.5
+            class_dict: {
+                0: "Neoplastic",
+                1: "Inflammatory",
+                2: "Connective",
+                3: "Dead",
+                4: "Epithelial",
+            }
+            postproc_tile_shape: [2048, 2048]
+    ioconfig:
+        class: io_config.IOSegmentorConfig
+        kwargs:
+            input_resolutions:
+                - {'units': 'mpp', 'resolution': 0.25}
+            output_resolutions:
+                - {'units': 'mpp', 'resolution': 0.25}
+            patch_input_shape: [256, 256]
+            patch_output_shape: [256, 256]
+            stride_shape: [240, 240]
+            save_resolution: {'units': 'baseline', 'resolution': 1.0}
+
+KongNet_MONKEY_1:
+    hf_repo_id: TIACentre/KongNet_pretrained_weights
+    architecture:
+        class: kongnet.KongNet
+        kwargs:
+            num_heads: 3
+            num_channels_per_head: [3, 3, 3]
+            target_channels: [2, 5, 8]
+            wide_decoder: True
+            min_distance: 11
+            threshold_abs: 0.5
+            class_dict: {
+                0: "Overall_Inflammatory",
+                1: "Lymphocyte",
+                2: "Monocyte",
+            }
+            postproc_tile_shape: [2048, 2048]
+    ioconfig:
+        class: io_config.IOSegmentorConfig
+        kwargs:
+            input_resolutions:
+                - {'units': 'mpp', 'resolution': 0.25}
+            output_resolutions:
+                - {'units': 'mpp', 'resolution': 0.25}
+            patch_input_shape: [256, 256]
+            patch_output_shape: [256, 256]
+            stride_shape: [224, 224]
+            save_resolution: {'units': 'baseline', 'resolution': 1.0}
+
+KongNet_PUMA_T1_3:
+    hf_repo_id: TIACentre/KongNet_pretrained_weights
+    architecture:
+        class: kongnet.KongNet
+        kwargs:
+            num_heads: 3
+            num_channels_per_head: [3, 3, 3]
+            target_channels: [2, 5, 8]
+            wide_decoder: False
+            min_distance: 13
+            threshold_abs: 0.5
+            class_dict: {
+                0: "Tumour_Cell",
+                1: "Lymphocyte",
+                2: "Other_Cell",
+            }
+            postproc_tile_shape: [2048, 2048]
+    ioconfig:
+        class: io_config.IOSegmentorConfig
+        kwargs:
+            input_resolutions:
+                - {'units': 'mpp', 'resolution': 0.25}
+            output_resolutions:
+                - {'units': 'mpp', 'resolution': 0.25}
+            patch_input_shape: [256, 256]
+            patch_output_shape: [256, 256]
+            stride_shape: [224, 224]
+            save_resolution: {'units': 'baseline', 'resolution': 1.0}
+
+KongNet_PUMA_T2_3:
+    hf_repo_id: TIACentre/KongNet_pretrained_weights
+    architecture:
+        class: kongnet.KongNet
+        kwargs:
+            num_heads: 10
+            num_channels_per_head: [3, 3, 3, 3, 3, 3, 3, 3, 3, 3]
+            target_channels: [2, 5, 8, 11, 14, 17, 20, 23, 26, 29]
+            wide_decoder: False
+            min_distance: 13
+            threshold_abs: 0.5
+            class_dict: {
+                0: "Tumour_Cell",
+                1: "Lymphocyte",
+                2: "Plasma_Cell",
+                3: "Histiocyte",
+                4: "Melanophage",
+                5: "Neutrophil",
+                6: "Stroma_Cell",
+                7: "Epithelial_Cell",
+                8: "Endothelial_Cell",
+                9: "Apoptotic_Cell",
+            }
+            postproc_tile_shape: [2048, 2048]
+    ioconfig:
+        class: io_config.IOSegmentorConfig
+        kwargs:
+            input_resolutions:
+                - {'units': 'mpp', 'resolution': 0.25}
+            output_resolutions:
+                - {'units': 'mpp', 'resolution': 0.25}
+            patch_input_shape: [256, 256]
+            patch_output_shape: [256, 256]
+            stride_shape: [224, 224]
+            save_resolution: {'units': 'baseline', 'resolution': 1.0}
+
+KongNet_Det_MIDOG_1:
+    hf_repo_id: TIACentre/KongNet_pretrained_weights
+    architecture:
+        class: kongnet.KongNet
+        kwargs:
+            num_heads: 1
+            num_channels_per_head: [1]
+            target_channels: [0]
+            wide_decoder: False
+            min_distance: 21
+            threshold_abs: 0.99
+            class_dict: {
+                0: "Mitotic_Figure"
+            }
+            postproc_tile_shape: [2048, 2048]
+    ioconfig:
+        class: io_config.IOSegmentorConfig
+        kwargs:
+            input_resolutions:
+                - {'units': 'mpp', 'resolution': 0.5}
+            output_resolutions:
+                - {'units': 'mpp', 'resolution': 0.5}
+            patch_input_shape: [512, 512]
+            patch_output_shape: [512, 512]
+            stride_shape: [492, 492]
+            save_resolution: {'units': 'baseline', 'resolution': 1.0}

tiatoolbox/models/architecture/grandqc.py

@@ -8,8 +8,6 @@
 
 Key Components:
 ---------------
-- SegmentationHead:
-    Final layer for segmentation output.
 - Conv2dReLU:
     Convolutional block with BatchNorm and ReLU activation.
 - DecoderBlock:
@@ -60,77 +58,10 @@
 from torch import nn
 
 from tiatoolbox.models.architecture.timm_efficientnet import EfficientNetEncoder
+from tiatoolbox.models.architecture.utils import SegmentationHead
 from tiatoolbox.models.models_abc import ModelABC
 
 
-class SegmentationHead(nn.Sequential):
-    """Segmentation head for UNet++ architecture.
-
-    This class defines the final segmentation layer for the UNet++ model.
-    It applies a convolution followed by optional upsampling and activation
-    to produce the segmentation output.
-
-    Attributes:
-        conv2d (nn.Conv2d):
-            Convolutional layer for feature transformation.
-        upsampling_layer (nn.Module):
-            Upsampling layer (bilinear interpolation or identity).
-        activation (nn.Module):
-            Activation function applied after upsampling.
-
-    Example:
-        >>> head = SegmentationHead(in_channels=64, out_channels=2)
-        >>> x = torch.randn(1, 64, 128, 128)
-        >>> output = head(x)
-        >>> output.shape
-        ... torch.Size([1, 2, 128, 128])
-
-    """
-
-    def __init__(
-        self: SegmentationHead,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int = 3,
-        activation: nn.Module | None = None,
-        upsampling: int = 1,
-    ) -> None:
-        """Initialize the SegmentationHead module.
-
-        This method sets up the segmentation head by creating a convolutional layer,
-        an optional upsampling layer, and an activation function. It is typically
-        used as the final stage in UNet++ architectures for semantic segmentation.
-
-        Args:
-            in_channels (int):
-                Number of input channels to the segmentation head.
-            out_channels (int):
-                Number of output channels (usually equal to the number of classes).
-            kernel_size (int):
-                Size of the convolution kernel. Defaults to 3.
-            activation (nn.Module | None):
-                Activation function applied after convolution. Defaults to None.
-            upsampling (int):
-                Upsampling factor applied to the output. Defaults to 1.
-
-        Raises:
-            ValueError:
-                If `kernel_size` or `upsampling` is not a positive integer.
-
-        """
-        conv2d = nn.Conv2d(
-            in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2
-        )
-        upsampling_layer = (
-            nn.UpsamplingBilinear2d(scale_factor=upsampling)
-            if upsampling > 1
-            else nn.Identity()
-        )
-        if activation is None:
-            activation = nn.Identity()
-        super().__init__(conv2d, upsampling_layer, activation)
-
-
 class Conv2dReLU(nn.Sequential):
     """Conv2d + BatchNorm + ReLU block.