address comments

Author: Jiaqi-Lv

tests/models/test_arch_grandqc.py

@@ -1,9 +1,13 @@
 """Unit test package for GrandQC Tissue Model."""
 
+from collections.abc import Callable
+from pathlib import Path
+
 import numpy as np
 import torch
 from torch import nn
 
+from tiatoolbox.annotation.storage import SQLiteStore
 from tiatoolbox.models.architecture import (
     fetch_pretrained_weights,
     get_pretrained_model,
@@ -15,16 +19,17 @@
     UnetPlusPlusDecoder,
 )
 from tiatoolbox.models.engine.io_config import IOSegmentorConfig
-from tiatoolbox.utils.misc import select_device
+from tiatoolbox.models.engine.semantic_segmentor import SemanticSegmentor
+from tiatoolbox.utils import env_detection as toolbox_env
 from tiatoolbox.wsicore.wsireader import VirtualWSIReader
 
-ON_GPU = False
+device = "cuda" if toolbox_env.has_gpu() else "cpu"
 
 
 def test_functional_grandqc() -> None:
     """Test for GrandQC model."""
     # test fetch pretrained weights
-    pretrained_weights = fetch_pretrained_weights("grandqc_tissue_detection_mpp10")
+    pretrained_weights = fetch_pretrained_weights("grandqc_tissue_detection")
     assert pretrained_weights is not None
 
     # test creation
@@ -36,7 +41,7 @@ def test_functional_grandqc() -> None:
     model.load_state_dict(pretrained)
 
     # test get pretrained model
-    model, ioconfig = get_pretrained_model("grandqc_tissue_detection_mpp10")
+    model, ioconfig = get_pretrained_model("grandqc_tissue_detection")
     assert isinstance(model, GrandQCModel)
     assert isinstance(ioconfig, IOSegmentorConfig)
     assert model.num_output_channels == 2
@@ -54,7 +59,7 @@ def test_functional_grandqc() -> None:
         ],
     )
     batch = torch.from_numpy(batch)
-    output = model.infer_batch(model, batch, device=select_device(on_gpu=ON_GPU))
+    output = model.infer_batch(model, batch, device=device)
     assert output.shape == (2, 512, 512, 2)
 
 
@@ -76,6 +81,39 @@ def test_grandqc_preproc_postproc() -> None:
     assert postproc_image.dtype == np.int64
 
 
+def test_grandqc_with_semantic_segmentor(
+    remote_sample: Callable, track_tmp_path: Path
+) -> None:
+    """Test GrandQC tissue mask generation."""
+    segmentor = SemanticSegmentor(model="grandqc_tissue_detection")
+
+    sample_image = remote_sample("svs-1-small")
+    inputs = [str(sample_image)]
+
+    output = segmentor.run(
+        images=inputs,
+        device=device,
+        patch_mode=False,
+        output_type="annotationstore",
+        save_dir=track_tmp_path / "grandqc_test_outputs",
+        overwrite=True,
+    )
+
+    assert len(output) == 1
+    assert Path(output[sample_image]).exists()
+
+    store = SQLiteStore.open(output[sample_image])
+    assert len(store) == 3
+
+    tissue_area_px = 0.0
+    for annotation in store.values():
+        assert annotation.properties["type"] == "mask"
+        tissue_area_px += annotation.geometry.area
+    assert 3003401 < tissue_area_px < 3003402
+
+    store.close()
+
+
 def test_segmentation_head_behaviour() -> None:
     """Verify SegmentationHead defaults and upsampling."""
     head = SegmentationHead(3, 5, activation=None, upsampling=1)

tiatoolbox/data/pretrained_model.yaml

@@ -935,7 +935,7 @@ nuclick_light-pannuke:
             patch_output_shape: [128, 128]
             save_resolution: {'units': 'baseline', 'resolution': 1.0}
 
-grandqc_tissue_detection_mpp10:
+grandqc_tissue_detection:
     hf_repo_id: TIACentre/GrandQC_Tissue_Detection
     architecture:
         class: grandqc.GrandQCModel

tiatoolbox/models/architecture/grandqc.py

@@ -1,7 +1,7 @@
 """GrandQC Tissue Detection Model Architecture.
 
 This module defines the GrandQC model for tissue detection in digital pathology.
-It implements a UNet++ architecture with an EfficientNet encoder and a segmentation
+It implements a UNet++ architecture with an EfficientNetB0 encoder and a segmentation
 head for high-resolution tissue segmentation. The model is designed to identify
 tissue regions and background areas for quality control in whole slide images (WSIs).
 
@@ -205,7 +205,7 @@ class DecoderBlock(nn.Module):
 
     This block performs upsampling and feature fusion using skip connections
     from the encoder. It consists of two convolutional layers with ReLU activation
-    and optional attention mechanisms.
+    and optional attention mechanisms (not implemented).
 
     Attributes:
         conv1 (Conv2dReLU):
@@ -222,9 +222,9 @@ class DecoderBlock(nn.Module):
 
     Example:
         >>> block = DecoderBlock(in_channels=128, skip_channels=64, out_channels=64)
-        >>> x = torch.randn(1, 128, 64, 64)
+        >>> input_tensor = torch.randn(1, 128, 64, 64)
         >>> skip = torch.randn(1, 64, 128, 128)
-        >>> output = block(x, skip)
+        >>> output = block(input_tensor, skip)
         >>> output.shape
         ... torch.Size([1, 64, 128, 128])
 
@@ -268,7 +268,7 @@ def __init__(
 
     def forward(
         self: DecoderBlock,
-        x: torch.Tensor,
+        input_tensor: torch.Tensor,
         skip: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Forward pass through the decoder block.
@@ -277,29 +277,33 @@ def forward(
         (if provided), and applies two convolutional layers with attention.
 
         Args:
-            x (torch.Tensor):
-                Input tensor from the previous decoder layer.
+            input_tensor (torch.Tensor):
+                (B, C_in, H, W). Input tensor from the previous decoder layer.
             skip (torch.Tensor | None):
+                (B, C_skip, H*2, W*2).
                 Skip connection tensor from the encoder. Defaults to None.
 
         Returns:
             torch.Tensor:
+                (B, C_out, H*2, W*2).
                 Output tensor after decoding and feature refinement.
 
         """
-        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        input_tensor = torch.nn.functional.interpolate(
+            input_tensor, scale_factor=2.0, mode="nearest"
+        )
         if skip is not None:
-            x = torch.cat([x, skip], dim=1)
-            x = self.attention1(x)
-        x = self.conv1(x)
-        x = self.conv2(x)
-        return self.attention2(x)
+            input_tensor = torch.cat([input_tensor, skip], dim=1)
+            input_tensor = self.attention1(input_tensor)
+        input_tensor = self.conv1(input_tensor)
+        input_tensor = self.conv2(input_tensor)
+        return self.attention2(input_tensor)
 
 
 class CenterBlock(nn.Sequential):
     """Center block for UNet++ architecture.
 
-    This block is placed at the bottleneck of the UNet++ architecture.
+    This block can be placed at the bottleneck of the UNet++ architecture.
     It consists of two convolutional layers with ReLU activation, used
     to process the deepest feature maps before decoding begins.
 
@@ -311,8 +315,8 @@ class CenterBlock(nn.Sequential):
 
     Example:
         >>> center = CenterBlock(in_channels=256, out_channels=512)
-        >>> x = torch.randn(1, 256, 32, 32)
-        >>> output = center(x)
+        >>> input_tensor = torch.randn(1, 256, 32, 32)
+        >>> output = center(input_tensor)
         >>> output.shape
         ... torch.Size([1, 512, 32, 32])