kaiko-eva 0.3.1__tar.gz → 0.3.3__tar.gz

{kaiko_eva-0.3.1 → kaiko_eva-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kaiko-eva
-Version: 0.3.1
+Version: 0.3.3
 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>

{kaiko_eva-0.3.1 → kaiko_eva-0.3.3}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "kaiko-eva"
-version = "0.3.1"
+version = "0.3.3"
 description = "Evaluation Framework for oncology foundation models."
 keywords = [
     "machine-learning",

{kaiko_eva-0.3.1 → kaiko_eva-0.3.3}/src/eva/core/data/dataloaders/dataloader.py

@@ -77,7 +77,9 @@ class DataLoader:
             shuffle=self.shuffle,
             sampler=sampler or self.sampler,
             batch_sampler=self.batch_sampler,
-            num_workers=self.num_workers or multiprocessing.cpu_count(),
+            num_workers=(
+                multiprocessing.cpu_count() if self.num_workers is None else self.num_workers
+            ),
             collate_fn=self.collate_fn,
             pin_memory=self.pin_memory,
             drop_last=self.drop_last,

{kaiko_eva-0.3.1 → kaiko_eva-0.3.3}/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.1 → kaiko_eva-0.3.3}/src/eva/core/loggers/log/__init__.py

@@ -2,5 +2,6 @@
 
 from eva.core.loggers.log.image import log_image
 from eva.core.loggers.log.parameters import log_parameters
+from eva.core.loggers.log.table import log_table
 
-__all__ = ["log_image", "log_parameters"]
+__all__ = ["log_image", "log_parameters", "log_table"]

kaiko_eva-0.3.3/src/eva/core/loggers/log/table.py

@@ -0,0 +1,73 @@
+"""Table log functionality."""
+
+import functools
+from typing import List
+
+import pandas as pd
+
+from eva.core.loggers import loggers
+from eva.core.loggers.log import utils
+
+
+@functools.singledispatch
+def log_table(
+    logger,
+    tag: str,
+    columns: List[str] | None = None,
+    data: List[List[str]] | None = None,
+    dataframe: pd.DataFrame | None = None,
+    step: int = 0,
+) -> None:
+    """Adds a a table to the logger.
+
+    The table can be defined either with `columns` and `data` or with `dataframe`.
+
+    Args:
+        logger: The logger to log the table to.
+        tag: The log tag.
+        columns: The column names of the table.
+        data: The data of the table as a list of lists.
+        dataframe: A pandas DataFrame to log.
+        step: The global step of the log.
+    """
+    utils.raise_not_supported(logger, "table")
+
+
+@log_table.register
+def _(
+    loggers: list,
+    tag: str,
+    columns: List[str] | None = None,
+    data: List[List[str]] | None = None,
+    dataframe: pd.DataFrame | None = None,
+    step: int = 0,
+) -> None:
+    """Adds a table to a list of supported loggers."""
+    for logger in loggers:
+        log_table(
+            logger,
+            tag=tag,
+            columns=columns,
+            data=data,
+            dataframe=dataframe,
+            step=step,
+        )
+
+
+@log_table.register
+def _(
+    logger: loggers.WandbLogger,
+    tag: str,
+    columns: List[str] | None = None,
+    data: List[List[str]] | None = None,
+    dataframe: pd.DataFrame | None = None,
+    step: int = 0,
+) -> None:
+    """Adds a table to a Wandb logger."""
+    logger.log_table(
+        key=tag,
+        columns=columns,
+        data=data,
+        dataframe=dataframe,
+        step=step,
+    )

{kaiko_eva-0.3.1 → kaiko_eva-0.3.3}/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.1 → kaiko_eva-0.3.3}/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.1 → kaiko_eva-0.3.3}/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.1 → kaiko_eva-0.3.3}/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.1 → kaiko_eva-0.3.3}/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.1 → kaiko_eva-0.3.3}/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.3/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.1 → kaiko_eva-0.3.3}/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",