kaiko-eva 0.3.0__tar.gz → 0.3.2__tar.gz

{kaiko_eva-0.3.0 → kaiko_eva-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kaiko-eva
-Version: 0.3.0
+Version: 0.3.2
 Summary: Evaluation Framework for oncology foundation models.
 Keywords: machine-learning,evaluation-framework,oncology,foundation-models
 Author-Email: Ioannis Gatopoulos <ioannis@kaiko.ai>, =?utf-8?q?Nicolas_K=C3=A4nzig?= <nicolas@kaiko.ai>, Roman Moser <roman@kaiko.ai>
@@ -243,7 +243,7 @@ Requires-Dist: scipy>=1.14.0; extra == "vision"
 Requires-Dist: monai>=1.3.2; extra == "vision"
 Requires-Dist: einops>=0.8.1; extra == "vision"
 Provides-Extra: language
-Requires-Dist: datasets>=3.2.0; extra == "language"
+Requires-Dist: datasets<4.0.0,>=2.19.0; extra == "language"
 Requires-Dist: litellm>=1.61.8; extra == "language"
 Provides-Extra: all
 Requires-Dist: h5py>=3.10.0; extra == "all"
@@ -258,7 +258,7 @@ Requires-Dist: imagesize>=1.4.1; extra == "all"
 Requires-Dist: scipy>=1.14.0; extra == "all"
 Requires-Dist: monai>=1.3.2; extra == "all"
 Requires-Dist: einops>=0.8.1; extra == "all"
-Requires-Dist: datasets>=3.2.0; extra == "all"
+Requires-Dist: datasets<4.0.0,>=2.19.0; extra == "all"
 Requires-Dist: litellm>=1.61.8; extra == "all"
 Description-Content-Type: text/markdown
 

{kaiko_eva-0.3.0 → kaiko_eva-0.3.2}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "kaiko-eva"
-version = "0.3.0"
+version = "0.3.2"
 description = "Evaluation Framework for oncology foundation models."
 keywords = [
     "machine-learning",
@@ -73,7 +73,7 @@ vision = [
     "einops>=0.8.1",
 ]
 language = [
-    "datasets>=3.2.0",
+    "datasets<4.0.0,>=2.19.0",
     "litellm>=1.61.8",
 ]
 all = [
@@ -89,7 +89,7 @@ all = [
     "scipy>=1.14.0",
     "monai>=1.3.2",
     "einops>=0.8.1",
-    "datasets>=3.2.0",
+    "datasets<4.0.0,>=2.19.0",
     "litellm>=1.61.8",
 ]
 

{kaiko_eva-0.3.0 → kaiko_eva-0.3.2}/src/eva/core/interface/interface.py

@@ -102,6 +102,9 @@ class Interface:
             model: The model module to use but not modify.
             data: The data module containing validation data.
         """
+        if getattr(data.datasets, "val", None) is None:
+            raise ValueError("The provided data module does not contain a validation dataset.")
+
         eva_trainer.run_evaluation_session(
             base_trainer=trainer,
             base_model=model,
@@ -110,3 +113,22 @@ class Interface:
             n_runs=trainer.n_runs,
             verbose=trainer.n_runs > 1,
         )
+
+    def test(
+        self,
+        trainer: eva_trainer.Trainer,
+        model: modules.ModelModule,
+        data: datamodules.DataModule,
+    ) -> None:
+        """Same as validate, but runs the test stage."""
+        if getattr(data.datasets, "test", None) is None:
+            raise ValueError("The provided data module does not contain a test dataset.")
+
+        eva_trainer.run_evaluation_session(
+            base_trainer=trainer,
+            base_model=model,
+            datamodule=data,
+            stages=["test"],
+            n_runs=trainer.n_runs,
+            verbose=trainer.n_runs > 1,
+        )

{kaiko_eva-0.3.0 → kaiko_eva-0.3.2}/src/eva/core/trainers/_recorder.py

@@ -86,10 +86,13 @@ class SessionRecorder:
 
     def update(
         self,
-        validation_scores: _EVALUATE_OUTPUT,
+        validation_scores: _EVALUATE_OUTPUT | None = None,
         test_scores: _EVALUATE_OUTPUT | None = None,
     ) -> None:
         """Updates the state of the tracked metrics in-place."""
+        if validation_scores is None and test_scores is None:
+            raise ValueError("At least one of validation_scores or test_scores must be provided.")
+
         self._update_validation_metrics(validation_scores)
         self._update_test_metrics(test_scores)
 
@@ -117,9 +120,10 @@ class SessionRecorder:
         self._validation_metrics = []
         self._test_metrics = []
 
-    def _update_validation_metrics(self, metrics: _EVALUATE_OUTPUT) -> None:
+    def _update_validation_metrics(self, metrics: _EVALUATE_OUTPUT | None) -> None:
         """Updates the validation metrics in-place."""
-        self._validation_metrics = _update_session_metrics(self._validation_metrics, metrics)
+        if metrics:
+            self._validation_metrics = _update_session_metrics(self._validation_metrics, metrics)
 
     def _update_test_metrics(self, metrics: _EVALUATE_OUTPUT | None) -> None:
         """Updates the test metrics in-place."""

{kaiko_eva-0.3.0 → kaiko_eva-0.3.2}/src/eva/core/trainers/functional.py

@@ -47,7 +47,7 @@ def run_evaluation_session(
             stages=stages,
             verbose=not verbose,
         )
-        if validation_scores:
+        if validation_scores or test_scores:
             recorder.update(validation_scores, test_scores)
     recorder.save()
 
@@ -89,7 +89,7 @@ def run_evaluation(
 
     if "fit" in stages:
         trainer.fit(model, datamodule=datamodule)
-    if "validate" in stages:
+    if "validate" in stages and getattr(datamodule.datasets, "val", None) is not None:
         validation_scores = trainer.validate(
             model=model,
             datamodule=datamodule,

{kaiko_eva-0.3.0 → kaiko_eva-0.3.2}/src/eva/vision/data/datasets/__init__.py

@@ -22,7 +22,6 @@ from eva.vision.data.datasets.segmentation import (
     LiTS17,
     MoNuSAC,
     MSDTask7Pancreas,
-    TotalSegmentator2D,
 )
 from eva.vision.data.datasets.vision import VisionDataset
 from eva.vision.data.datasets.wsi import MultiWsiDataset, WsiDataset
@@ -40,7 +39,6 @@ __all__ = [
     "PANDASmall",
     "Camelyon16",
     "PatchCamelyon",
-    "TotalSegmentator2D",
     "UniToPatho",
     "WsiClassificationDataset",
     "CoNSeP",

{kaiko_eva-0.3.0 → kaiko_eva-0.3.2}/src/eva/vision/data/datasets/segmentation/__init__.py

@@ -7,7 +7,6 @@ from eva.vision.data.datasets.segmentation.embeddings import EmbeddingsSegmentat
 from eva.vision.data.datasets.segmentation.lits17 import LiTS17
 from eva.vision.data.datasets.segmentation.monusac import MoNuSAC
 from eva.vision.data.datasets.segmentation.msd_task7_pancreas import MSDTask7Pancreas
-from eva.vision.data.datasets.segmentation.total_segmentator_2d import TotalSegmentator2D
 
 __all__ = [
     "BCSS",
@@ -17,5 +16,4 @@ __all__ = [
     "LiTS17",
     "MSDTask7Pancreas",
     "MoNuSAC",
-    "TotalSegmentator2D",
 ]

{kaiko_eva-0.3.0 → kaiko_eva-0.3.2}/src/eva/vision/models/networks/decoders/segmentation/__init__.py

@@ -9,6 +9,7 @@ from eva.vision.models.networks.decoders.segmentation.semantic import (
     ConvDecoderWithImage,
     SingleLinearDecoder,
     SwinUNETRDecoder,
+    SwinUNETRDecoderWithProjection,
 )
 
 __all__ = [
@@ -20,4 +21,5 @@ __all__ = [
     "LinearDecoder",
     "SingleLinearDecoder",
     "SwinUNETRDecoder",
+    "SwinUNETRDecoderWithProjection",
 ]

{kaiko_eva-0.3.0 → kaiko_eva-0.3.2}/src/eva/vision/models/networks/decoders/segmentation/semantic/__init__.py

@@ -5,7 +5,10 @@ from eva.vision.models.networks.decoders.segmentation.semantic.common import (
     ConvDecoderMS,
     SingleLinearDecoder,
 )
-from eva.vision.models.networks.decoders.segmentation.semantic.swin_unetr import SwinUNETRDecoder
+from eva.vision.models.networks.decoders.segmentation.semantic.swin_unetr import (
+    SwinUNETRDecoder,
+    SwinUNETRDecoderWithProjection,
+)
 from eva.vision.models.networks.decoders.segmentation.semantic.with_image import (
     ConvDecoderWithImage,
 )
@@ -16,4 +19,5 @@ __all__ = [
     "ConvDecoderWithImage",
     "SingleLinearDecoder",
     "SwinUNETRDecoder",
+    "SwinUNETRDecoderWithProjection",
 ]

kaiko_eva-0.3.2/src/eva/vision/models/networks/decoders/segmentation/semantic/swin_unetr.py

@@ -0,0 +1,255 @@
+"""Decoder based on Swin UNETR."""
+
+from typing import List
+
+import torch
+from monai.networks.blocks import dynunet_block, unetr_block
+from torch import nn
+
+
+class SwinUNETRDecoder(nn.Module):
+    """Swin transformer decoder based on UNETR [0].
+
+    - [0] UNETR: Transformers for 3D Medical Image Segmentation
+      https://arxiv.org/pdf/2103.10504
+    """
+
+    def __init__(
+        self,
+        out_channels: int,
+        feature_size: int = 48,
+        spatial_dims: int = 3,
+    ) -> None:
+        """Builds the decoder.
+
+        Args:
+            out_channels: Number of output channels.
+            feature_size: Dimension of network feature size.
+            spatial_dims: Number of spatial dimensions.
+        """
+        super().__init__()
+
+        self.decoder5 = unetr_block.UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=16 * feature_size,
+            out_channels=8 * feature_size,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.decoder4 = unetr_block.UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=feature_size * 8,
+            out_channels=feature_size * 4,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.decoder3 = unetr_block.UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=feature_size * 4,
+            out_channels=feature_size * 2,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.decoder2 = unetr_block.UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=feature_size * 2,
+            out_channels=feature_size,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.decoder1 = unetr_block.UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=feature_size,
+            out_channels=feature_size,
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.out = dynunet_block.UnetOutBlock(
+            spatial_dims=spatial_dims,
+            in_channels=feature_size,
+            out_channels=out_channels,
+        )
+
+    def _forward_features(self, features: List[torch.Tensor]) -> torch.Tensor:
+        """Forward function for multi-level feature maps to a single one."""
+        enc0, enc1, enc2, enc3, hid3, dec4 = features
+        dec3 = self.decoder5(dec4, hid3)
+        dec2 = self.decoder4(dec3, enc3)
+        dec1 = self.decoder3(dec2, enc2)
+        dec0 = self.decoder2(dec1, enc1)
+        out = self.decoder1(dec0, enc0)
+        return self.out(out)
+
+    def forward(self, features: List[torch.Tensor]) -> torch.Tensor:
+        """Maps the patch embeddings to a segmentation mask.
+
+        Args:
+            features: List of multi-level intermediate features from
+                :class:`SwinUNETREncoder`.
+
+        Returns:
+            Tensor containing scores for all of the classes with shape
+            (batch_size, n_classes, image_height, image_width).
+        """
+        return self._forward_features(features)
+
+
+class SwinUNETRDecoderWithProjection(nn.Module):
+    """Swin transformer decoder based on UNETR [0].
+
+    This implementation adds additional projection layers to reduce
+    the number of channels in the feature maps before applying the upscaling
+    convolutional blocks. This reduces the number of trainable parameters
+    significantly and is useful when scaling up the encoder architecture.
+
+    - [0] UNETR: Transformers for 3D Medical Image Segmentation
+      https://arxiv.org/pdf/2103.10504
+    """
+
+    def __init__(
+        self,
+        out_channels: int,
+        feature_size: int = 48,
+        spatial_dims: int = 3,
+        project_dims: list[int] | None = None,
+        checkpoint_path: str | None = None,
+    ) -> None:
+        """Builds the decoder.
+
+        Args:
+            out_channels: Number of output channels.
+            feature_size: Dimension of network feature size.
+            spatial_dims: Number of spatial dimensions.
+            project_dims: List of 6 dimensions to project encoder features to.
+                If None, uses default channel progression based on feature_size.
+                This is not part of the original implementation, but helps
+                to reduce the number of decoder parameters when scaling up
+                the encoder architecture (feature_size).
+            checkpoint_path: Path to the checkpoint file.
+        """
+        super().__init__()
+
+        self._checkpoint_path = checkpoint_path
+        self._project_dims = project_dims
+
+        if project_dims is not None and len(project_dims) != 6:
+            raise ValueError(
+                f"project_dims must have exactly 6 dimensions, got {len(project_dims)}"
+            )
+
+        channel_dims = project_dims or [
+            feature_size,
+            feature_size,
+            feature_size * 2,
+            feature_size * 4,
+            feature_size * 8,
+            feature_size * 16,
+        ]
+
+        self.decoder5 = unetr_block.UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=channel_dims[5],
+            out_channels=channel_dims[4],
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.decoder4 = unetr_block.UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=channel_dims[4],
+            out_channels=channel_dims[3],
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.decoder3 = unetr_block.UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=channel_dims[3],
+            out_channels=channel_dims[2],
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.decoder2 = unetr_block.UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=channel_dims[2],
+            out_channels=channel_dims[1],
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.decoder1 = unetr_block.UnetrUpBlock(
+            spatial_dims=spatial_dims,
+            in_channels=channel_dims[1],
+            out_channels=channel_dims[0],
+            kernel_size=3,
+            upsample_kernel_size=2,
+            norm_name="instance",
+            res_block=True,
+        )
+        self.out = dynunet_block.UnetOutBlock(
+            spatial_dims=spatial_dims,
+            in_channels=channel_dims[0],
+            out_channels=out_channels,
+        )
+
+        if self._project_dims:
+            conv_layer = nn.Conv2d if spatial_dims == 2 else nn.Conv3d
+            self.proj_enc0 = conv_layer(feature_size, self._project_dims[0], kernel_size=1)
+            self.proj_enc1 = conv_layer(feature_size, self._project_dims[1], kernel_size=1)
+            self.proj_enc2 = conv_layer(feature_size * 2, self._project_dims[2], kernel_size=1)
+            self.proj_enc3 = conv_layer(feature_size * 4, self._project_dims[3], kernel_size=1)
+            self.proj_hid3 = conv_layer(feature_size * 8, self._project_dims[4], kernel_size=1)
+            self.proj_dec4 = conv_layer(feature_size * 16, self._project_dims[5], kernel_size=1)
+
+    def _forward_features(self, features: List[torch.Tensor]) -> torch.Tensor:
+        """Forward function for multi-level feature maps to a single one."""
+        enc0, enc1, enc2, enc3, hid3, dec4 = self._project_features(features)
+        dec3 = self.decoder5(dec4, hid3)
+        dec2 = self.decoder4(dec3, enc3)
+        dec1 = self.decoder3(dec2, enc2)
+        dec0 = self.decoder2(dec1, enc1)
+        out = self.decoder1(dec0, enc0)
+        return self.out(out)
+
+    def _project_features(self, features: List[torch.Tensor]) -> List[torch.Tensor]:
+        """Projects features using 1x1 to reduce number of channels."""
+        return (
+            [
+                self.proj_enc0(features[0]),
+                self.proj_enc1(features[1]),
+                self.proj_enc2(features[2]),
+                self.proj_enc3(features[3]),
+                self.proj_hid3(features[4]),
+                self.proj_dec4(features[5]),
+            ]
+            if self._project_dims
+            else features
+        )
+
+    def forward(self, features: List[torch.Tensor]) -> torch.Tensor:
+        """Maps the patch embeddings to a segmentation mask.
+
+        Args:
+            features: List of multi-level intermediate features from
+                :class:`SwinUNETREncoder`.
+
+        Returns:
+            Tensor containing scores for all of the classes with shape
+            (batch_size, n_classes, image_height, image_width).
+        """
+        return self._forward_features(features)

{kaiko_eva-0.3.0 → kaiko_eva-0.3.2}/tests/eva/core/models/wrappers/test_from_torchub.py

@@ -8,6 +8,7 @@ import torch
 from eva.core.models import wrappers
 
 
+@pytest.mark.skip(reason="Test triggers rate limits")
 @pytest.mark.parametrize(
     "model_name, repo_or_dir, out_indices, model_kwargs, "
     "input_tensor, expected_len, expected_shape",

kaiko_eva-0.3.2/tests/eva/vision/models/networks/decoders/segmentation/test_swin_unetr.py

@@ -0,0 +1,198 @@
+"""Unit test for the SwinUNETRDecoder."""
+
+from typing import List
+
+import pytest
+import torch
+
+from eva.vision.models.networks.backbones.radiology.swin_unetr import SwinUNETREncoder
+from eva.vision.models.networks.decoders.segmentation.semantic import (
+    SwinUNETRDecoder,
+    SwinUNETRDecoderWithProjection,
+)
+
+
+@pytest.mark.parametrize("spatial_dims", [2, 3])
+def test_forward_pass(spatial_dims: int):
+    """Tests if the full forward pass using generated features by encoder."""
+    in_channels, num_classes, feature_size, out_indices = 1, 7, 3, 6
+    spatial_dim_size = (96, 96, 96) if spatial_dims == 3 else (96, 96)
+    batch = torch.randn(4, in_channels, *spatial_dim_size)
+
+    encoder = SwinUNETREncoder(
+        in_channels=in_channels,
+        feature_size=feature_size,
+        spatial_dims=spatial_dims,
+        out_indices=out_indices,
+    )
+    decoder = SwinUNETRDecoder(
+        out_channels=num_classes,
+        feature_size=feature_size,
+        spatial_dims=spatial_dims,
+    )
+    features = encoder(batch)
+    assert len(features) == out_indices
+
+    output = decoder(features)
+    assert output.shape == torch.Size([4, num_classes, *spatial_dim_size])
+
+
+@pytest.mark.parametrize("spatial_dims", [2, 3])
+@pytest.mark.parametrize("feature_size", [3, 6])
+@pytest.mark.parametrize("out_channels", [1, 2, 5])
+def test_swin_unetr_decoder_shapes(spatial_dims: int, feature_size: int, out_channels: int):
+    """Tests SwinUNETRDecoder with different configurations."""
+    batch_size = 2
+    spatial_dim_size = (32, 32, 32) if spatial_dims == 3 else (32, 32)
+
+    decoder = SwinUNETRDecoder(
+        out_channels=out_channels,
+        feature_size=feature_size,
+        spatial_dims=spatial_dims,
+    )
+
+    features = _create_mock_features(batch_size, feature_size, spatial_dim_size)
+    output = decoder(features)
+    assert output.shape == torch.Size([batch_size, out_channels, *spatial_dim_size])
+
+
+@pytest.mark.parametrize("spatial_dims", [2, 3])
+@pytest.mark.parametrize("feature_size", [3, 6])
+@pytest.mark.parametrize("out_channels", [1, 2])
+def test_swin_unetr_decoder_with_projection_shapes(
+    spatial_dims: int, feature_size: int, out_channels: int
+):
+    """Tests SwinUNETRDecoderWithProjection with different configurations."""
+    batch_size = 2
+    spatial_dim_size = (32, 32, 32) if spatial_dims == 3 else (32, 32)
+
+    decoder = SwinUNETRDecoderWithProjection(
+        out_channels=out_channels,
+        feature_size=feature_size,
+        spatial_dims=spatial_dims,
+    )
+
+    features = _create_mock_features(batch_size, feature_size, spatial_dim_size)
+    output = decoder(features)
+    assert output.shape == torch.Size([batch_size, out_channels, *spatial_dim_size])
+
+
+@pytest.mark.parametrize("spatial_dims", [2, 3])
+def test_swin_unetr_decoder_with_projection_custom_dims(spatial_dims: int):
+    """Tests SwinUNETRDecoderWithProjection with custom projection dimensions."""
+    batch_size = 2
+    feature_size = 6
+    out_channels = 2
+    spatial_dim_size = (32, 32, 32) if spatial_dims == 3 else (32, 32)
+    project_dims = [3, 4, 6, 8, 12, 16]
+
+    decoder = SwinUNETRDecoderWithProjection(
+        out_channels=out_channels,
+        feature_size=feature_size,
+        spatial_dims=spatial_dims,
+        project_dims=project_dims,
+    )
+
+    features = _create_mock_features(batch_size, feature_size, spatial_dim_size)
+    output = decoder(features)
+    assert output.shape == torch.Size([batch_size, out_channels, *spatial_dim_size])
+
+
+def test_swin_unetr_decoder_with_projection_invalid_project_dims():
+    """Tests SwinUNETRDecoderWithProjection with invalid projection dimensions."""
+    with pytest.raises(ValueError, match="project_dims must have exactly 6 dimensions"):
+        SwinUNETRDecoderWithProjection(
+            out_channels=2,
+            feature_size=6,
+            spatial_dims=3,
+            project_dims=[3, 4, 6, 8],  # Only 4 dims instead of 6
+        )
+
+
+def test_swin_unetr_decoder_deterministic():
+    """Tests that the decoder produces deterministic outputs."""
+    torch.manual_seed(42)
+
+    batch_size = 1
+    feature_size = 3
+    out_channels = 2
+    spatial_dims = 3
+    spatial_dim_size = (64, 64, 64)
+
+    decoder = SwinUNETRDecoder(
+        out_channels=out_channels,
+        feature_size=feature_size,
+        spatial_dims=spatial_dims,
+    )
+
+    features = _create_mock_features(batch_size, feature_size, spatial_dim_size)
+
+    output1 = decoder(features)
+    output2 = decoder(features)
+
+    assert torch.allclose(output1, output2, atol=1e-6)
+
+
+def test_swin_unetr_decoder_with_projection_no_projection():
+    """Tests SwinUNETRDecoderWithProjection without projection dimensions."""
+    batch_size = 1
+    feature_size = 3
+    out_channels = 2
+    spatial_dims = 3
+    spatial_dim_size = (64, 64, 64)
+
+    decoder = SwinUNETRDecoderWithProjection(
+        out_channels=out_channels,
+        feature_size=feature_size,
+        spatial_dims=spatial_dims,
+        project_dims=None,  # No projection
+    )
+
+    features = _create_mock_features(batch_size, feature_size, spatial_dim_size)
+    output = decoder(features)
+    assert output.shape == torch.Size([batch_size, out_channels, *spatial_dim_size])
+
+
+def test_swin_unetr_decoder_wrong_features_length():
+    """Tests decoder behavior with wrong number of features."""
+    decoder = SwinUNETRDecoder(
+        out_channels=2,
+        feature_size=3,
+        spatial_dims=3,
+    )
+
+    # Provide only 4 features instead of 6
+    features = [
+        torch.randn(1, 3, 16, 16, 16),
+        torch.randn(1, 3, 8, 8, 8),
+        torch.randn(1, 6, 4, 4, 4),
+        torch.randn(1, 12, 2, 2, 2),
+    ]
+
+    with pytest.raises((ValueError, RuntimeError, IndexError)):
+        decoder(features)
+
+
+def _create_mock_features(
+    batch_size: int,
+    feature_size: int,
+    spatial_dim_size: tuple[int, ...],
+) -> List[torch.Tensor]:
+    """Create mock features matching expected SwinUNETREncoder output.
+
+    Args:
+        batch_size: Batch size for the features.
+        feature_size: Base feature size.
+        spatial_dim_size: Spatial dimensions of the input.
+
+    Returns:
+        List of 6 feature tensors: [enc0, enc1, enc2, enc3, hid3, dec4].
+    """
+    return [
+        torch.randn(batch_size, feature_size, *spatial_dim_size),  # enc0
+        torch.randn(batch_size, feature_size, *(s // 2 for s in spatial_dim_size)),  # enc1
+        torch.randn(batch_size, feature_size * 2, *(s // 4 for s in spatial_dim_size)),  # enc2
+        torch.randn(batch_size, feature_size * 4, *(s // 8 for s in spatial_dim_size)),  # enc3
+        torch.randn(batch_size, feature_size * 8, *(s // 16 for s in spatial_dim_size)),  # hid3
+        torch.randn(batch_size, feature_size * 16, *(s // 32 for s in spatial_dim_size)),  # dec4
+    ]