remove smp dependency

Author: Jiaqi-Lv

requirements/requirements.txt

@@ -28,7 +28,6 @@ requests>=2.28.1
 scikit-image>=0.20
 scikit-learn>=1.2.0
 scipy>=1.8
-segmentation-models-pytorch>=0.5.0
 shapely>=2.0.0
 SimpleITK>=2.2.1
 sphinx>=5.3.0

tests/models/test_arch_grandqc.py

@@ -7,7 +7,7 @@
     fetch_pretrained_weights,
     get_pretrained_model,
 )
-from tiatoolbox.models.architecture.grandqc import TissueDetectionModel
+from tiatoolbox.models.architecture.grandqc import GrandQCModel
 from tiatoolbox.models.engine.io_config import IOSegmentorConfig
 from tiatoolbox.utils.misc import select_device
 from tiatoolbox.wsicore.wsireader import VirtualWSIReader
@@ -22,7 +22,7 @@ def test_functional_grandqc() -> None:
     assert pretrained_weights is not None
 
     # test creation
-    model = TissueDetectionModel(num_input_channels=3, num_output_channels=2)
+    model = GrandQCModel(num_input_channels=3, num_output_channels=2)
     assert model is not None
 
     # load pretrained weights
@@ -31,7 +31,7 @@ def test_functional_grandqc() -> None:
 
     # test get pretrained model
     model, ioconfig = get_pretrained_model("grandqc_tissue_detection_mpp10")
-    assert isinstance(model, TissueDetectionModel)
+    assert isinstance(model, GrandQCModel)
     assert isinstance(ioconfig, IOSegmentorConfig)
     assert model.num_input_channels == 3
     assert model.num_output_channels == 2
@@ -54,7 +54,7 @@ def test_functional_grandqc() -> None:
 
 def test_grandqc_preproc_postproc() -> None:
     """Test GrandQC preproc and postproc functions."""
-    model = TissueDetectionModel(num_input_channels=3, num_output_channels=2)
+    model = GrandQCModel(num_input_channels=3, num_output_channels=2)
 
     generator = np.random.default_rng(1337)
     # test preproc

tiatoolbox/data/pretrained_model.yaml

@@ -938,7 +938,7 @@ nuclick_light-pannuke:
 grandqc_tissue_detection_mpp10:
     hf_repo_id: TIACentre/GrandQC_Tissue_Detection
     architecture:
-        class: grandqc.TissueDetectionModel
+        class: grandqc.GrandQCModel
         kwargs:
             num_input_channels: 3
             num_output_channels: 2

tiatoolbox/models/architecture/grandqc.py

@@ -9,14 +9,14 @@
 
 import cv2
 import numpy as np
-import segmentation_models_pytorch as smp
 import torch
 
+from tiatoolbox.models.architecture.unetplusplus import UNetPlusPlusModel
 from tiatoolbox.models.models_abc import ModelABC
 
 
-class TissueDetectionModel(ModelABC):
-    """GrandQC Tissue Detection Model.
+class GrandQCModel(ModelABC):
+    """GrandQC Tissue Detection Model [1].
 
     Example:
         >>> from tiatoolbox.models.engine.semantic_segmentor import SemanticSegmentor
@@ -32,28 +32,39 @@ class TissueDetectionModel(ModelABC):
         ...     output_type="annotationstore",
         ... )
 
+    References:
+        [1] Weng Z. et al. "GrandQC: a comprehensive solution to quality control problem
+        in digital pathology".
+        Nature Communications 2024
+
     """
 
     def __init__(
-        self: TissueDetectionModel, num_input_channels: int, num_output_channels: int
+        self: GrandQCModel, num_input_channels: int, num_output_channels: int
     ) -> None:
         """Initialize TissueDetectionModel."""
         super().__init__()
         self.num_input_channels = num_input_channels
         self.num_output_channels = num_output_channels
         self._postproc = self.postproc
         self._preproc = self.preproc
-        self.tissue_detection_model = smp.UnetPlusPlus(
-            encoder_name="timm-efficientnet-b0",
-            encoder_weights=None,
-            in_channels=self.num_input_channels,
+        self.tissue_detection_model = UNetPlusPlusModel(
             classes=self.num_output_channels,
-            activation=None,
         )
 
     @staticmethod
     def preproc(image: np.ndarray) -> np.ndarray:
-        """Apply jpg compression then ImageNet normalise."""
+        """Apply JPEG compression and ImageNet normalization to the input image.
+
+        Args:
+            image (np.ndarray):
+                Input image as a NumPy array (H, W, C) in uint8 format.
+
+        Returns:
+            np.ndarray:
+                The preprocessed image.
+
+        """
         encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 80]
         _, compressed_image = cv2.imencode(".jpg", image, encode_param)
         compressed_image = np.array(cv2.imdecode(compressed_image, 1))
@@ -69,11 +80,19 @@ def postproc(image: np.ndarray) -> np.ndarray:
         This simply applies argmin to obtain tissue class.
         (Tissue = 0, Background = 1)
 
+        Args:
+            image (np.ndarray):
+                Input probability map as a NumPy array (H, W, C).
+
+        Returns:
+            np.ndarray:
+                Tissue mask
+
         """
         return image.argmin(axis=-1)
 
     def forward(
-        self: TissueDetectionModel,
+        self: GrandQCModel,
         imgs: torch.Tensor,
         *args: tuple[Any, ...],  # skipcq: PYL-W0613  # noqa: ARG002
         **kwargs: dict,  # skipcq: PYL-W0613  # noqa: ARG002
@@ -120,9 +139,9 @@ def infer_batch(
         return probs.cpu().numpy()
 
     def load_state_dict(
-        self: TissueDetectionModel,
+        self: GrandQCModel,
         state_dict: Mapping[str, Any],
         **kwargs: bool,
     ) -> torch.nn.modules.module._IncompatibleKeys:
-        """Load state dict for the TissueDetectionModel."""
+        """Load state dict for the GrandQCModel."""
         return self.tissue_detection_model.load_state_dict(state_dict, **kwargs)