🆕 Add Virchow2 distilled model configs

tests/models/test_arch_segformer.py

@@ -164,6 +164,37 @@ def test_segformer_forward_and_metadata() -> None:
     assert output.shape == (1, 2, 64, 64)
 
 
+def test_segformer_forward_matches_huggingface_module() -> None:
+    """Direct stage execution should match the Hugging Face segmentation head."""
+    model = Segformer(encoder_name="mit_b0", classes=2)
+    model.eval()
+    x = torch.randn(1, 3, 64, 64)
+
+    with torch.inference_mode():
+        output = model(x)
+        reference = model.activation(
+            model.upsampling(model.model(pixel_values=x).logits)
+        )
+
+    assert torch.allclose(output, reference, atol=1e-5)
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices")
+def test_segformer_dataparallel_forward() -> None:
+    """DataParallel inference should keep per-stage feature maps aligned."""
+    model = Segformer(encoder_name="mit_b0", classes=2)
+    model.eval()
+    wrapped = model.to("cuda")
+    batch = torch.randn(2, 3, 64, 64)
+
+    with torch.inference_mode():
+        single = model(batch.cuda())
+        parallel = wrapped(batch.cuda())
+
+    assert parallel.shape == (2, 2, 64, 64)
+    assert torch.allclose(parallel.cpu(), single.cpu(), atol=1e-4)
+
+
 def test_segformer_preproc_imagenet_normalization() -> None:
     """Preprocessing should apply ImageNet channel normalization."""
     image = np.ones((8, 8, 3), dtype=np.uint8) * 128
@@ -187,10 +218,14 @@ def test_segformer_postproc_with_numpy_and_dask() -> None:
     probs_da = da.from_array(probs_np, chunks=(16, 16, 2))
     mask_da = model.postproc(probs_da)
 
+    assert isinstance(mask_np, np.ndarray)
     assert mask_np.shape == (32, 32)
     assert mask_np.dtype == np.uint8
+
+    assert isinstance(mask_da, da.Array)
     assert mask_da.shape == (32, 32)
     assert mask_da.dtype == np.uint8
+    np.testing.assert_array_equal(mask_da.compute(), mask_np)
 
 
 def test_segformer_infer_batch_probability_output() -> None:
@@ -310,3 +345,20 @@ def test_native_hf_state_dict_still_loads() -> None:
     incompatible = other.load_state_dict(state, strict=True)
     assert incompatible.missing_keys == []
     assert incompatible.unexpected_keys == []
+
+
+def test_wrapped_model_key_checkpoint_loads() -> None:
+    """Checkpoints nested under a ``model`` key should unwrap and load."""
+    model = Segformer(
+        encoder_name="mit_b0", classes=3, decoder_segmentation_channels=32
+    )
+    wrapped = {"model": model.state_dict(), "epoch": 10}
+
+    other = Segformer(
+        encoder_name="mit_b0",
+        classes=3,
+        decoder_segmentation_channels=32,
+    )
+    incompatible = other.load_state_dict(wrapped, strict=True)
+    assert incompatible.missing_keys == []
+    assert incompatible.unexpected_keys == []

tests/test_utils.py

@@ -2433,6 +2433,51 @@ def test_semantic_segmentation_returns_json_dict() -> None:
         assert feature["class_value"] in class_dict
 
 
+def test_semantic_segmentation_qupath_json_skips_empty_classes() -> None:
+    """Classes with no pixels should be skipped when contour hierarchy is None."""
+    preds_np = np.ones((8, 8), dtype=np.uint8)
+
+    qupath_json = _semantic_segmentations_as_qupath_json(
+        layer_list=[1, 2],
+        preds=da.from_array(preds_np, chunks=(8, 8)),
+        scale_factor=(1.0, 1.0),
+        class_dict={1: "Stroma", 2: "Tumour"},
+        save_path=None,
+        verbose=False,
+    )
+
+    assert qupath_json["type"] == "FeatureCollection"
+    assert len(qupath_json["features"]) == 1
+    assert qupath_json["features"][0]["class_value"] == 1
+    assert qupath_json["features"][0]["name"] == "Stroma"
+
+
+def test_semantic_segmentation_qupath_json_preserves_holes() -> None:
+    """Nested tissue classes should become holes in the surrounding polygon."""
+    preds_np = np.ones((16, 16), dtype=np.uint8)
+    preds_np[5:11, 5:11] = 2
+
+    qupath_json = _semantic_segmentations_as_qupath_json(
+        layer_list=[1, 2],
+        preds=da.from_array(preds_np, chunks=(16, 16)),
+        scale_factor=(1.0, 1.0),
+        class_dict={1: "Stroma", 2: "Tumour"},
+        save_path=None,
+        verbose=False,
+    )
+
+    stroma = [feat for feat in qupath_json["features"] if feat["class_value"] == 1]
+    tumour = [feat for feat in qupath_json["features"] if feat["class_value"] == 2]
+
+    assert len(stroma) == 1
+    assert len(tumour) == 1
+    stroma_geom = stroma[0]["geometry"]
+    assert stroma_geom["type"] == "Polygon"
+    assert len(stroma_geom["coordinates"]) == 2
+    assert tumour[0]["geometry"]["type"] == "Polygon"
+    assert len(tumour[0]["geometry"]["coordinates"]) == 1
+
+
 def test_dict_to_store_patch_predictions_returns_qupath_json() -> None:
     """Test for dict_to_store_patch_predictions QuPath JSON dict."""
     # Fake patch output

tiatoolbox/data/pretrained_model.yaml

@@ -1189,3 +1189,192 @@ efficientunet-tissue_mask:
             stride_shape: [480, 480]
             save_resolution: {'units': 'mpp', 'resolution': 8.0}
             ignore_index: 0
+
+
+segformerb5_distilled_PUMA_1:
+    hf_repo_id: TIACentre/Virchow2_Distilled_Segmentation_Models
+    catalog:
+        description: >
+            Segment Tumour, Stroma, Necrosis, Blood Vessels, Epidermis in melanoma.
+        primary site: skin
+        supported_sites:
+            - skin
+    architecture:
+        class: segformer.Segformer
+        kwargs:
+            encoder_name: "mit_b5"
+            decoder_segmentation_channels: 768
+            in_channels: 3
+            classes: 6
+            class_dict: {
+                0: "Background",
+                1: "Stroma",
+                2: "Blood Vessels",
+                3: "Tumour",
+                4: "Epidermis",
+                5: "Necrosis",
+            }
+    ioconfig:
+        class: io_config.IOSegmentorConfig
+        kwargs:
+            input_resolutions:
+                - {'units': 'mpp', 'resolution': 1.14}
+            output_resolutions:
+                - {'units': 'mpp', 'resolution': 1.14}
+            patch_input_shape: [224, 224]
+            patch_output_shape: [224, 224]
+            stride_shape: [112, 112]
+            save_resolution: {'units': 'mpp', 'resolution': 1.14}
+            ignore_index: 0
+
+
+segformerb5_distilled_PUMA_2:
+    hf_repo_id: TIACentre/Virchow2_Distilled_Segmentation_Models
+    catalog:
+        description: >
+            Segment Tumour, Stroma, Necrosis, Blood Vessels, Epidermis in melanoma.
+        primary site: skin
+        supported_sites:
+            - skin
+    architecture:
+        class: segformer.Segformer
+        kwargs:
+            encoder_name: "mit_b5"
+            decoder_segmentation_channels: 768
+            in_channels: 3
+            classes: 6
+            class_dict: {
+                0: "Background",
+                1: "Stroma",
+                2: "Blood Vessels",
+                3: "Tumour",
+                4: "Epidermis",
+                5: "Necrosis",
+            }
+    ioconfig:
+        class: io_config.IOSegmentorConfig
+        kwargs:
+            input_resolutions:
+                - {'units': 'mpp', 'resolution': 1.14}
+            output_resolutions:
+                - {'units': 'mpp', 'resolution': 1.14}
+            patch_input_shape: [224, 224]
+            patch_output_shape: [224, 224]
+            stride_shape: [112, 112]
+            save_resolution: {'units': 'mpp', 'resolution': 1.14}
+            ignore_index: 0
+
+
+segformerb5_distilled_PUMA_5:
+    hf_repo_id: TIACentre/Virchow2_Distilled_Segmentation_Models
+    catalog:
+        description: >
+            Segment Tumour, Stroma, Necrosis, Blood Vessels, Epidermis in melanoma.
+        primary site: skin
+        supported_sites:
+            - skin
+    architecture:
+        class: segformer.Segformer
+        kwargs:
+            encoder_name: "mit_b5"
+            decoder_segmentation_channels: 768
+            in_channels: 3
+            classes: 6
+            class_dict: {
+                0: "Background",
+                1: "Stroma",
+                2: "Blood Vessels",
+                3: "Tumour",
+                4: "Epidermis",
+                5: "Necrosis",
+            }
+    ioconfig:
+        class: io_config.IOSegmentorConfig
+        kwargs:
+            input_resolutions:
+                - {'units': 'mpp', 'resolution': 1.14}
+            output_resolutions:
+                - {'units': 'mpp', 'resolution': 1.14}
+            patch_input_shape: [224, 224]
+            patch_output_shape: [224, 224]
+            stride_shape: [112, 112]
+            save_resolution: {'units': 'mpp', 'resolution': 1.14}
+            ignore_index: 0
+
+
+segformerb5_distilled_IGNITE_1:
+    hf_repo_id: TIACentre/Virchow2_Distilled_Segmentation_Models
+    catalog:
+        description: >
+            Segment Tumour, Stroma, Necrosis, Macrophages, Rest in non-small cell lung cancer and its metastatic sites.
+        primary site: lung
+        supported_sites:
+            - lung
+            - liver
+            - brain
+            - adrenal gland
+            - kidney
+    architecture:
+        class: segformer.Segformer
+        kwargs:
+            encoder_name: "mit_b5"
+            decoder_segmentation_channels: 768
+            in_channels: 3
+            classes: 7
+            class_dict: {
+                0: "Unknown",
+                1: "Background",
+                2: "Tumour",
+                3: "Stroma",
+                4: "Necrosis",
+                5: "Macrophages",
+                6: "Rest",
+            }
+    ioconfig:
+        class: io_config.IOSegmentorConfig
+        kwargs:
+            input_resolutions:
+                - {'units': 'mpp', 'resolution': 0.5}
+            output_resolutions:
+                - {'units': 'mpp', 'resolution': 0.5}
+            patch_input_shape: [448, 448]
+            patch_output_shape: [448, 448]
+            stride_shape: [224, 224]
+            save_resolution: {'units': 'mpp', 'resolution': 0.5}
+            ignore_index: 0
+
+
+segformerb5_distilled_BEETLE_1:
+    hf_repo_id: TIACentre/Virchow2_Distilled_Segmentation_Models
+    catalog:
+        description: >
+            Segment Invasive epithelium, Non-invasive epithelium, Necrosis in breast cancer images.
+        primary site: breast
+        supported_sites:
+            - breast
+    architecture:
+        class: segformer.Segformer
+        kwargs:
+            encoder_name: "mit_b5"
+            decoder_segmentation_channels: 768
+            in_channels: 3
+            classes: 5
+            class_dict: {
+                0: "Unknown",
+                1: "Other",
+                2: "Non_invasive_epithelium",
+                3: "Invasive_epithelium",
+                4: "Necrosis",
+            }
+    ioconfig:
+        class: io_config.IOSegmentorConfig
+        kwargs:
+            input_resolutions:
+                - {'units': 'mpp', 'resolution': 1.0}
+            output_resolutions:
+                - {'units': 'mpp', 'resolution': 1.0}
+            patch_input_shape: [448, 448]
+            patch_output_shape: [448, 448]
+            stride_shape: [224, 224]
+            save_resolution: {'units': 'mpp', 'resolution': 1.0}
+            ignore_index: 0