initial update

Author: Jiaqi-Lv

examples/nucleus-detection.ipynb

@@ -0,0 +1,166 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fa584486",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "|2025-12-17|18:50:54.872| [WARNING] /home/u1910100/miniconda3/envs/tiatoolbox-dev/lib/python3.12/site-packages/albumentations/__init__.py:28: UserWarning: A new version of Albumentations is available: '2.0.8' (you have '2.0.4'). Upgrade using: pip install -U albumentations. To disable automatic update checks, set the environment variable NO_ALBUMENTATIONS_UPDATE to 1.\n",
+      "  check_for_updates()\n",
+      "\n"
+     ]
+    },
+    {
+     "ename": "",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
+      "\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
+      "\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
+      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
+     ]
+    }
+   ],
+   "source": [
+    "import sys\n",
+    "sys.path.append(\"../\")\n",
+    "import pathlib\n",
+    "from tiatoolbox.models.engine.nucleus_detector import NucleusDetector\n",
+    "from tiatoolbox.models.architecture.kongnet import KongNet"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "78f694c7",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "|2025-12-17|18:50:57.923| [WARNING] GPU is not compatible with torch.compile. Compatible GPUs include NVIDIA V100, A100, and H100. Speedup numbers may be lower than expected.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "|2025-12-17|18:50:58.126| [INFO] When providing multiple whole slide images, the outputs will be saved and the locations of outputs will be returned to the calling function when `run()` finishes successfully.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c6892643281f4d72844868ee40e1375d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Processing WSIs:   0%|          | 0/1 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "|2025-12-17|18:50:58.311| [WARNING] Read: Scale > 1.This means that the desired resolution is higher than the WSI baseline (maximum encoded resolution). Interpolation of read regions may occur.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3342b87637924ddcaa77bb121f27f4e7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Inferring patches:   0%|          | 0/1004 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Current Memory usage: 75.52868374109333 %  exceeds specified threshold: 75. Saving intermediate results to disk.\n"
+     ]
+    }
+   ],
+   "source": [
+    "detector = NucleusDetector(model='KongNet_CoNIC_1')\n",
+    "\n",
+    "wsi_path = \"/media/u1910100/data/slides/TUM1.svs\"\n",
+    "\n",
+    "out = detector.run(\n",
+    "    images=[pathlib.Path(wsi_path)],\n",
+    "    patch_mode=False,\n",
+    "    device=\"cuda\",\n",
+    "    save_dir=pathlib.Path(\"/media/u1910100/data/overlays/test\"),\n",
+    "    overwrite=True,\n",
+    "    output_type=\"annotationstore\",\n",
+    "    auto_get_mask=True,\n",
+    "    memory_threshold=75,\n",
+    "    num_workers=1,\n",
+    "    batch_size=8,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b559e6d0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from tiatoolbox.annotation.storage import SQLiteStore\n",
+    "\n",
+    "store_path = \"/media/u1910100/data/overlays/test/wsi1_2k_2k.db\"\n",
+    "store = SQLiteStore.open(store_path)\n",
+    "\n",
+    "for ann in store.values():\n",
+    "    print(ann)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e0db521a",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "tiatoolbox-dev",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

test.py

@@ -0,0 +1,24 @@
+
+import pathlib
+from tiatoolbox.models.engine.nucleus_detector import NucleusDetector
+
+
+detector = NucleusDetector(model='KongNet_CoNIC_1')
+
+# wsi_path = "/media/u1910100/data/slides/TCGA-AO-A0J2-01Z-00-DX1.7C9FEC7B-6040-4C58-9563-D10C0D7AC72E.svs"
+wsi_path = "/media/u1910100/data/slides/CMU-1-Small-Region.svs"
+
+out = detector.run(
+    images=[pathlib.Path(wsi_path)],
+    patch_mode=False,
+    device="cuda",
+    save_dir=pathlib.Path("/media/u1910100/data/overlays/test"),
+    overwrite=True,
+    output_type="annotationstore",
+    auto_get_mask=True,
+    memory_threshold=30,
+    num_workers=1,
+    batch_size=8,
+    cache_dir=pathlib.Path("/media/u1910100/data/cache")
+)
+

tests/models/test_arch_kongnet.py

@@ -0,0 +1,50 @@
+"""Unit test package for KongNet Model."""
+
+from collections.abc import Callable
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from torch import nn
+
+from tiatoolbox.annotation.storage import SQLiteStore
+from tiatoolbox.models.architecture import (
+    fetch_pretrained_weights,
+    get_pretrained_model,
+)
+from tiatoolbox.models.architecture.kongnet import KongNet
+from tiatoolbox.utils import env_detection as toolbox_env
+from tiatoolbox.wsicore.wsireader import VirtualWSIReader
+
+device = "cuda" if toolbox_env.has_gpu() else "cpu"
+
+
+def test_KongNet_Modeling() -> None:
+    """Test for KongNet model."""
+
+    # test creation
+    model = KongNet(
+        num_heads=6, 
+        num_channels_per_head=[3,3,3,3,3,3],
+        target_channels=[2, 5, 8, 11, 14, 17],
+        min_distance=5,
+        threshold_abs=0.5,
+        wide_decoder=False,
+    )
+
+    # ckp_path = "/media/u1910100/data/Monkey/conic_models/efficientnetv2_l_eq_loss/KongNet_CoNIC_1.pth"
+    # state_dict = torch.load(ckp_path, map_location=device)["model"]
+    # model.load_state_dict(state_dict)
+    model = model.to(device)
+    # assert model is not None
+
+    model.eval()
+    with torch.no_grad():
+        input_tensor = torch.randn(1, 3, 256, 256).to(device)
+        output = model(input_tensor)
+        assert output.shape == (1, 6, 256, 256)
+
+        batch_tensor = torch.randn(4, 3, 256, 256).to(device)
+        output = model(batch_tensor)
+        assert output.shape == (4, 6, 256, 256)

tests/models/test_arch_utils.py

@@ -7,9 +7,11 @@
 
 from tiatoolbox.models.architecture.utils import (
     UpSample2x,
+    Attention,
     centre_crop,
     centre_crop_to_shape,
     peak_detection_map_overlap,
+    nms_on_detection_maps,
 )
 
 
@@ -160,3 +162,40 @@ 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 = Attention(name=None, in_channels=16)
+    output = attention(test_input)
+    assert torch.sum(output - test_input) == 0
+
+    attention = Attention(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"):
+        attention = Attention(name="random_name", in_channels=16)
+

tiatoolbox/data/pretrained_model.yaml

@@ -964,3 +964,36 @@ 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}
+