fix tests

Author: Jiaqi-Lv

tests/models/test_arch_grandqc.py

@@ -25,7 +25,7 @@ def test_functional_grandqc(remote_sample: Callable) -> None:
     assert pretrained_weights is not None
 
     # test creation
-    model = TissueDetectionModel()
+    model = TissueDetectionModel(num_input_channels=3, num_output_channels=2)
     assert model is not None
 
     # load pretrained weights
@@ -36,6 +36,8 @@ def test_functional_grandqc(remote_sample: Callable) -> None:
     model, ioconfig = get_pretrained_model("grandqc_tissue_detection_mpp10")
     assert isinstance(model, TissueDetectionModel)
     assert isinstance(ioconfig, IOSegmentorConfig)
+    assert model.num_input_channels == 3
+    assert model.num_output_channels == 2
 
     # test inference
     mini_wsi_svs = Path(remote_sample("wsi2_4k_4k_svs"))

tiatoolbox/models/architecture/grandqc.py

@@ -34,15 +34,20 @@ class TissueDetectionModel(ModelABC):
 
     """
 
-    def __init__(self: TissueDetectionModel) -> None:
+    def __init__(
+        self: TissueDetectionModel, 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,
-            classes=2,
+            in_channels=self.num_input_channels,
+            classes=self.num_output_channels,
             activation=None,
         )