kaiko-eva 0.3.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

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,
+        )

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."""

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,

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",

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",
 ]

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",
 ]

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",
 ]

eva/vision/models/networks/decoders/segmentation/semantic/swin_unetr.py

@@ -102,3 +102,154 @@ class SwinUNETRDecoder(nn.Module):
             (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.dist-info → kaiko_eva-0.3.2.dist-info}/METADATA

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

{kaiko_eva-0.3.1.dist-info → kaiko_eva-0.3.2.dist-info}/RECORD

@@ -48,7 +48,7 @@ eva/core/data/transforms/padding/pad_2d_tensor.py,sha256=J4maGFmeQf9IHRxt5kU-6eI
 eva/core/data/transforms/sampling/__init__.py,sha256=BFKbvRjlZrwS0GcNrM54ZSWt6PrQARfFlXM1jJ-wpvo,149
 eva/core/data/transforms/sampling/sample_from_axis.py,sha256=Zbhp94lVa70WQKmSOKMTsOMe2c7wLqNZto7JqWhSdtI,1229
 eva/core/interface/__init__.py,sha256=chdpKXipxe1NP-Fgr_d9r6X1gMna0XiEa38waJ6FzTM,98
-eva/core/interface/interface.py,sha256=VNagHKsr2T7Ufm1rpA0NCnMi6F2qpKEmMBq_3eGZIRE,3826
+eva/core/interface/interface.py,sha256=EeBrIqUlO497BOZGZrtO-fZnc_BhIJrrqyQmRfqWXcw,4591
 eva/core/loggers/__init__.py,sha256=4YMLNlN9LnuKqhBI1R1keh69dmMD-2lcH3HKwwyn380,266
 eva/core/loggers/dummy.py,sha256=Y7ypH0ecSAIkkZ5LzTmNNEzlKkqeaHfUNMCDKVOg6D4,1204
 eva/core/loggers/experimental_loggers.py,sha256=p5uCK_9QCYufRhE-LZQUJWbhGElyobX_zRM78yX4p2o,230
@@ -99,9 +99,9 @@ eva/core/models/wrappers/huggingface.py,sha256=-_fA81YRnoMc7O7SbrnCEj1dM_xArpQ8W
 eva/core/models/wrappers/onnx.py,sha256=34li_xSwPryN8nJDrFyif_Hve1AEH7Ry9E_lZmf7JJM,1834
 eva/core/trainers/__init__.py,sha256=jhsKJF7HAae7EOiG3gKIAHH_h3dZlTE2JRcCHJmOzJc,208
 eva/core/trainers/_logging.py,sha256=gi4FqPy2GuVmh0WZY6mYwF7zMPvnoFA050B0XdCP6PU,2571
-eva/core/trainers/_recorder.py,sha256=uD17l_WVveFuWuann59VU9iJ-Jumdh9F6vnAcL3M_FU,7855
+eva/core/trainers/_recorder.py,sha256=M-BJHLgqGxR_MSV6f_WC7GN2JHYEEinV1-hNLpH667A,8062
 eva/core/trainers/_utils.py,sha256=M3h8lVhUmkeSiEXpX9hRdMvThGFCnTP15gv-hd1CZkc,321
-eva/core/trainers/functional.py,sha256=tsBfpXjEQ8BiBJ9wZWp0AUUOOxy7UUrLX4GSjQZTeCs,4510
+eva/core/trainers/functional.py,sha256=_Mw-NIPU2tPffxpK5t3sHBmVI6u163phCpoJFiauH7E,4583
 eva/core/trainers/trainer.py,sha256=a3OwLWOZKDqxayrd0ugUmxJKyQx6XDb4GHtdL8-AEV0,4826
 eva/core/utils/__init__.py,sha256=cndVBvtYxEW7hykH39GCNVI86zkXNn8Lw2A0sUJHS04,237
 eva/core/utils/clone.py,sha256=qcThZOuAs1cs0uV3BL5eKeM2VIBjuRPBe1t-NiUFM5Y,569
@@ -146,7 +146,7 @@ eva/vision/data/dataloaders/__init__.py,sha256=9ykBD4vyZ-Yv3IEnqvVcSMURS-gXWjOun
 eva/vision/data/dataloaders/collate_fn/__init__.py,sha256=GCvJaeILmAc_-lhGw8yzj2cC2KG4i1PvSWAyVzPKvVo,146
 eva/vision/data/dataloaders/collate_fn/collection.py,sha256=45s9fKjVBnqfnuGWmJZMtt_DDGnfuf7qkWe0QmxXMKo,611
 eva/vision/data/dataloaders/worker_init.py,sha256=lFWywHGCC4QxHeDXrneF8DQ45XG3WmVltEELJrPyLz0,1182
-eva/vision/data/datasets/__init__.py,sha256=s3h4w71LiM6dT6AYWzCG2-nexkSuuTWixw4KrCGAhS8,1026
+eva/vision/data/datasets/__init__.py,sha256=_04LqKv46oUXdmQAlSmWkgZYueHFwM-0iiOSMuFnFDw,976
 eva/vision/data/datasets/_utils.py,sha256=epPcaYE4w2_LtUKLLQJh6qQxUNVBe22JA06k4WUerYQ,1430
 eva/vision/data/datasets/_validators.py,sha256=77WZj8ewsuxUjW5WegJ-7zDuR6WdF5JbaOYdywhKIK4,2594
 eva/vision/data/datasets/classification/__init__.py,sha256=5fOGZxKGPeMCf3Jd9qAOYADPrkZnYg97_QE4DC79AMI,1074
@@ -161,7 +161,7 @@ eva/vision/data/datasets/classification/panda.py,sha256=HVfCvByyajdo5o_waqTpzZWC
 eva/vision/data/datasets/classification/patch_camelyon.py,sha256=1yXkfP680qxkQUFAPKRFbZv0cHAFx23s2vvT9th2nKM,7149
 eva/vision/data/datasets/classification/unitopatho.py,sha256=IO3msEsuOnmdcYZxF-eBpo0K97y54rWFmCb_KxuF4bk,5129
 eva/vision/data/datasets/classification/wsi.py,sha256=YMGxU8ECjudizt_uXUevuPS8k66HxtEQ7M2IZJmL6kE,4079
-eva/vision/data/datasets/segmentation/__init__.py,sha256=y_BjUj6kF-WeouSz0CCpPdOdX7n5hUrqsZGF68Xu9Hw,784
+eva/vision/data/datasets/segmentation/__init__.py,sha256=f0q9tzk4ahaZfrw_SgIE_puk_D7qmkSCKX1FP9aJITU,668
 eva/vision/data/datasets/segmentation/_utils.py,sha256=aXUHrnbefP6-OgSvDQHqssFKhUwETul_8aosqYiOfm8,3065
 eva/vision/data/datasets/segmentation/bcss.py,sha256=rqk6VqK0QCHLFnMnDuHd1JPJVK5_C6WnsmnNSKBw6Uo,8230
 eva/vision/data/datasets/segmentation/btcv.py,sha256=9rlEqGyb2SGJBY6Oj42FlHajQF8csf1Jq6jeuPSsfXI,8396
@@ -170,10 +170,8 @@ eva/vision/data/datasets/segmentation/embeddings.py,sha256=RsTuAwGEJPnWPY7q3pwcj
 eva/vision/data/datasets/segmentation/lits17.py,sha256=kcSCKxsgtUuCD1YEYvrb_L_BgOtZC8xDq1lX8ldSZc4,7635
 eva/vision/data/datasets/segmentation/metadata/__init__.py,sha256=o9Od0v6N9dNdf8hfefn2QaNNCD2sZMvc2K58zHA_Nrg,24
 eva/vision/data/datasets/segmentation/metadata/_msd_task7_pancreas.py,sha256=O2-ye0A7wIjcI_D857uvpYw-jckTqfhBUrhinqSNWq0,2553
-eva/vision/data/datasets/segmentation/metadata/_total_segmentator.py,sha256=DTaQaAisY7j1h0-zYk1_81Sr4b3D9PTMieYX0PMPtIc,3127
 eva/vision/data/datasets/segmentation/monusac.py,sha256=iv9-MFaTsGfGV1u6_lQNcSEeSpmVBDQC1Oa123iEtu0,8410
 eva/vision/data/datasets/segmentation/msd_task7_pancreas.py,sha256=dTsPD73PAP15VOXdHnX4eQqbpz2jGpCB31YISzinUd4,8964
-eva/vision/data/datasets/segmentation/total_segmentator_2d.py,sha256=TGz67AGuv8_Bm5DM5TyCtzRTuGXOuctZZNxdQtBxF1g,16987
 eva/vision/data/datasets/structs.py,sha256=RaTDW-B36PumcR5gymhCiX-r8GiKqIFcjqoEEjjFyUE,389
 eva/vision/data/datasets/vision.py,sha256=-_WRiyICMgqABR6Ay_RKBMfsPGwgx9MQfCA7WChHo24,3219
 eva/vision/data/datasets/wsi.py,sha256=dEAT_Si_Qb3qdSovUPeoiWeoPb7m-NGYqq44e3UXHk8,8384
@@ -256,13 +254,13 @@ eva/vision/models/networks/backbones/timm/backbones.py,sha256=ZbF9MMiL4Ylyy79XLe
 eva/vision/models/networks/backbones/universal/__init__.py,sha256=xgn3crSqlmUPYz-t2CR1zDKxhlyAEeApA-a6Y_eWQvc,417
 eva/vision/models/networks/backbones/universal/vit.py,sha256=To0OzwpuX5Y5PwjGidwV0Ssq3xa81dve081buwG_Ofg,3658
 eva/vision/models/networks/decoders/__init__.py,sha256=RXFWmoYw2i6E9VOUCJmU8c72icHannVuo-cUKy6fnLM,200
-eva/vision/models/networks/decoders/segmentation/__init__.py,sha256=SqmxtzxwBRF8g2hsiqe0o3Nr0HFK97azTnWLyqsYigY,652
+eva/vision/models/networks/decoders/segmentation/__init__.py,sha256=yVrRo2OisNRAlxDjWJGwipKA9HGeqRXd1ZL88eltoy4,726
 eva/vision/models/networks/decoders/segmentation/base.py,sha256=b2TIJKiJR9vejVRpNyedMJLPTrpHhAEXvco8atb9TPU,411
 eva/vision/models/networks/decoders/segmentation/decoder2d.py,sha256=HRonYTSriiq13aZCSNiYUc484qfOhkVT0yFiMW06CDc,4472
 eva/vision/models/networks/decoders/segmentation/linear.py,sha256=ui3-Y0rl4VEF75-sUghaF29P9wpxCVlp5iR_Ym-utUE,4666
-eva/vision/models/networks/decoders/segmentation/semantic/__init__.py,sha256=2yol7W1ARXL-Ge7gYxjUzaGTjH6nfMBlNqQJHprEWGg,539
+eva/vision/models/networks/decoders/segmentation/semantic/__init__.py,sha256=9QnepLMzQVE-wAZJXx0napVutg1HtkbDERcPsoevWGg,622
 eva/vision/models/networks/decoders/segmentation/semantic/common.py,sha256=FSf_eI-FaBroxPRJd4TiV97RCreauJh1IznIVzBT2eg,2528
-eva/vision/models/networks/decoders/segmentation/semantic/swin_unetr.py,sha256=ODUpnJrpDQl0m8CC2SPnE_lpFflzS0GSiCZOmrjL6uQ,3373
+eva/vision/models/networks/decoders/segmentation/semantic/swin_unetr.py,sha256=eSFvHng2lrc-Wd4g9CW4z8-yfKndbl0c7-sKhOautBU,9170
 eva/vision/models/networks/decoders/segmentation/semantic/with_image.py,sha256=I5PyGKKo8DcXYcw4xlCFzuavRJNRrzGT-szpDidMPXI,3516
 eva/vision/models/networks/decoders/segmentation/typings.py,sha256=rY4CXp0MNF16SHnx9TgGjXI_r8bVGSqAWdR835hXndg,537
 eva/vision/models/wrappers/__init__.py,sha256=ogmr-eeVuGaOCcsuxSp6PGyauP2QqWTb8dGTtbC7lRU,210
@@ -277,8 +275,8 @@ eva/vision/utils/io/image.py,sha256=IdOkr5MYqhYHz8U9drZ7wULTM3YHwCWSjZlu_Qdl4GQ,
 eva/vision/utils/io/mat.py,sha256=qpGifyjmpE0Xhv567Si7-zxKrgkgE0sywP70cHiLFGU,808
 eva/vision/utils/io/nifti.py,sha256=TFMgNhLqIK3sl3RjIRXEABM7FmSQjqVOwk1vXkuvX2w,4983
 eva/vision/utils/io/text.py,sha256=qYgfo_ZaDZWfG02NkVVYzo5QFySqdCCz5uLA9d-zXtI,701
-kaiko_eva-0.3.1.dist-info/METADATA,sha256=gXYGvp6Ap95944atE7L9Dxk8AnmuVhn22sHAC2iIl_g,25704
-kaiko_eva-0.3.1.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-kaiko_eva-0.3.1.dist-info/entry_points.txt,sha256=6CSLu9bmQYJSXEg8gbOzRhxH0AGs75BB-vPm3VvfcNE,88
-kaiko_eva-0.3.1.dist-info/licenses/LICENSE,sha256=e6AEzr7j_R-PYr2qLO-JwLn8y70jbVD3U2mxbRmwcI4,11338
-kaiko_eva-0.3.1.dist-info/RECORD,,
+kaiko_eva-0.3.2.dist-info/METADATA,sha256=3OdB75bdgEKDkAhIh75c3WcpevnOsemYcMfBEG0MKy8,25704
+kaiko_eva-0.3.2.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+kaiko_eva-0.3.2.dist-info/entry_points.txt,sha256=6CSLu9bmQYJSXEg8gbOzRhxH0AGs75BB-vPm3VvfcNE,88
+kaiko_eva-0.3.2.dist-info/licenses/LICENSE,sha256=e6AEzr7j_R-PYr2qLO-JwLn8y70jbVD3U2mxbRmwcI4,11338
+kaiko_eva-0.3.2.dist-info/RECORD,,

eva/vision/data/datasets/segmentation/metadata/_total_segmentator.py

@@ -1,91 +0,0 @@
-"""Utils for TotalSegmentator dataset classes."""
-
-from typing import Dict
-
-reduced_class_mappings: Dict[str, str] = {
-    # Abdominal Organs
-    "spleen": "spleen",
-    "kidney_right": "kidney",
-    "kidney_left": "kidney",
-    "gallbladder": "gallbladder",
-    "liver": "liver",
-    "stomach": "stomach",
-    "pancreas": "pancreas",
-    "small_bowel": "small_bowel",
-    "duodenum": "duodenum",
-    "colon": "colon",
-    # Endocrine System
-    "adrenal_gland_right": "adrenal_gland",
-    "adrenal_gland_left": "adrenal_gland",
-    "thyroid_gland": "thyroid_gland",
-    # Respiratory System
-    "lung_upper_lobe_left": "lungs",
-    "lung_lower_lobe_left": "lungs",
-    "lung_upper_lobe_right": "lungs",
-    "lung_middle_lobe_right": "lungs",
-    "lung_lower_lobe_right": "lungs",
-    "trachea": "trachea",
-    "esophagus": "esophagus",
-    # Urogenital System
-    "urinary_bladder": "urogenital_system",
-    "prostate": "urogenital_system",
-    "kidney_cyst_left": "kidney_cyst",
-    "kidney_cyst_right": "kidney_cyst",
-    # Vertebral Column
-    **{f"vertebrae_{v}": "vertebrae" for v in ["C1", "C2", "C3", "C4", "C5", "C6", "C7"]},
-    **{f"vertebrae_{v}": "vertebrae" for v in [f"T{i}" for i in range(1, 13)]},
-    **{f"vertebrae_{v}": "vertebrae" for v in [f"L{i}" for i in range(1, 6)]},
-    "vertebrae_S1": "vertebrae",
-    "sacrum": "sacral_spine",
-    # Cardiovascular System
-    "heart": "heart",
-    "aorta": "aorta",
-    "pulmonary_vein": "veins",
-    "brachiocephalic_trunk": "arteries",
-    "subclavian_artery_right": "arteries",
-    "subclavian_artery_left": "arteries",
-    "common_carotid_artery_right": "arteries",
-    "common_carotid_artery_left": "arteries",
-    "brachiocephalic_vein_left": "veins",
-    "brachiocephalic_vein_right": "veins",
-    "atrial_appendage_left": "atrial_appendage",
-    "superior_vena_cava": "veins",
-    "inferior_vena_cava": "veins",
-    "portal_vein_and_splenic_vein": "veins",
-    "iliac_artery_left": "arteries",
-    "iliac_artery_right": "arteries",
-    "iliac_vena_left": "veins",
-    "iliac_vena_right": "veins",
-    # Upper Extremity Bones
-    "humerus_left": "humerus",
-    "humerus_right": "humerus",
-    "scapula_left": "scapula",
-    "scapula_right": "scapula",
-    "clavicula_left": "clavicula",
-    "clavicula_right": "clavicula",
-    # Lower Extremity Bones
-    "femur_left": "femur",
-    "femur_right": "femur",
-    "hip_left": "hip",
-    "hip_right": "hip",
-    # Muscles
-    "gluteus_maximus_left": "gluteus",
-    "gluteus_maximus_right": "gluteus",
-    "gluteus_medius_left": "gluteus",
-    "gluteus_medius_right": "gluteus",
-    "gluteus_minimus_left": "gluteus",
-    "gluteus_minimus_right": "gluteus",
-    "autochthon_left": "autochthon",
-    "autochthon_right": "autochthon",
-    "iliopsoas_left": "iliopsoas",
-    "iliopsoas_right": "iliopsoas",
-    # Central Nervous System
-    "brain": "brain",
-    "spinal_cord": "spinal_cord",
-    # Skull and Thoracic Cage
-    "skull": "skull",
-    **{f"rib_left_{i}": "ribs" for i in range(1, 13)},
-    **{f"rib_right_{i}": "ribs" for i in range(1, 13)},
-    "costal_cartilages": "ribs",
-    "sternum": "sternum",
-}

eva/vision/data/datasets/segmentation/total_segmentator_2d.py

@@ -1,414 +0,0 @@
-"""TotalSegmentator 2D segmentation dataset class."""
-
-import functools
-import hashlib
-import os
-import re
-from glob import glob
-from pathlib import Path
-from typing import Any, Callable, Dict, List, Literal, Tuple
-
-import numpy as np
-import numpy.typing as npt
-import torch
-from torchvision import tv_tensors
-from torchvision.datasets import utils
-from typing_extensions import override
-
-from eva.core.utils import io as core_io
-from eva.core.utils import multiprocessing
-from eva.vision.data.datasets import _validators, structs, vision
-from eva.vision.data.datasets.segmentation.metadata import _total_segmentator
-from eva.vision.utils import io
-
-
-class TotalSegmentator2D(vision.VisionDataset[tv_tensors.Image, tv_tensors.Mask]):
-    """TotalSegmentator 2D segmentation dataset."""
-
-    _expected_dataset_lengths: Dict[str, int] = {
-        "train_small": 35089,
-        "val_small": 1283,
-        "train_full": 278190,
-        "val_full": 14095,
-        "test_full": 25578,
-    }
-    """Dataset version and split to the expected size."""
-
-    _sample_every_n_slices: int | None = None
-    """The amount of slices to sub-sample per 3D CT scan image."""
-
-    _resources_full: List[structs.DownloadResource] = [
-        structs.DownloadResource(
-            filename="Totalsegmentator_dataset_v201.zip",
-            url="https://zenodo.org/records/10047292/files/Totalsegmentator_dataset_v201.zip",
-            md5="fe250e5718e0a3b5df4c4ea9d58a62fe",
-        ),
-    ]
-    """Resources for the full dataset version."""
-
-    _resources_small: List[structs.DownloadResource] = [
-        structs.DownloadResource(
-            filename="Totalsegmentator_dataset_small_v201.zip",
-            url="https://zenodo.org/records/10047263/files/Totalsegmentator_dataset_small_v201.zip",
-            md5="6b5524af4b15e6ba06ef2d700c0c73e0",
-        ),
-    ]
-    """Resources for the small dataset version."""
-
-    _license: str = (
-        "Creative Commons Attribution 4.0 International "
-        "(https://creativecommons.org/licenses/by/4.0/deed.en)"
-    )
-    """Dataset license."""
-
-    def __init__(
-        self,
-        root: str,
-        split: Literal["train", "val", "test"] | None,
-        version: Literal["small", "full"] | None = "full",
-        download: bool = False,
-        classes: List[str] | None = None,
-        class_mappings: Dict[str, str] | None = _total_segmentator.reduced_class_mappings,
-        optimize_mask_loading: bool = True,
-        decompress: bool = True,
-        num_workers: int = 10,
-        transforms: Callable | None = None,
-    ) -> None:
-        """Initialize dataset.
-
-        Args:
-            root: Path to the root directory of the dataset. The dataset will
-                be downloaded and extracted here, if it does not already exist.
-            split: Dataset split to use. If `None`, the entire dataset is used.
-            version: The version of the dataset to initialize. If `None`, it will
-                use the files located at root as is and wont perform any checks.
-            download: Whether to download the data for the specified split.
-                Note that the download will be executed only by additionally
-                calling the :meth:`prepare_data` method and if the data does not
-                exist yet on disk.
-            classes: Whether to configure the dataset with a subset of classes.
-                If `None`, it will use all of them.
-            class_mappings: A dictionary that maps the original class names to a
-                reduced set of classes. If `None`, it will use the original classes.
-            optimize_mask_loading: Whether to pre-process the segmentation masks
-                in order to optimize the loading time. In the `setup` method, it
-                will reformat the binary one-hot masks to a semantic mask and store
-                it on disk.
-            decompress: Whether to decompress the ct.nii.gz files when preparing the data.
-                The label masks won't be decompressed, but when enabling optimize_mask_loading
-                it will export the semantic label masks to a single file in uncompressed .nii
-                format.
-            num_workers: The number of workers to use for optimizing the masks &
-                decompressing the .gz files.
-            transforms: A function/transforms that takes in an image and a target
-                mask and returns the transformed versions of both.
-
-        """
-        super().__init__(transforms=transforms)
-
-        self._root = root
-        self._split = split
-        self._version = version
-        self._download = download
-        self._classes = classes
-        self._optimize_mask_loading = optimize_mask_loading
-        self._decompress = decompress
-        self._num_workers = num_workers
-        self._class_mappings = class_mappings
-
-        if self._classes and self._class_mappings:
-            raise ValueError("Both 'classes' and 'class_mappings' cannot be set at the same time.")
-
-        self._samples_dirs: List[str] = []
-        self._indices: List[Tuple[int, int]] = []
-
-    @functools.cached_property
-    @override
-    def classes(self) -> List[str]:
-        def get_filename(path: str) -> str:
-            """Returns the filename from the full path."""
-            return os.path.basename(path).split(".")[0]
-
-        first_sample_labels = os.path.join(self._root, "s0011", "segmentations", "*.nii.gz")
-        all_classes = sorted(map(get_filename, glob(first_sample_labels)))
-        if self._classes:
-            is_subset = all(name in all_classes for name in self._classes)
-            if not is_subset:
-                raise ValueError("Provided class names are not subset of the original ones.")
-            classes = sorted(self._classes)
-        elif self._class_mappings:
-            is_subset = all(name in all_classes for name in self._class_mappings.keys())
-            if not is_subset:
-                raise ValueError("Provided class names are not subset of the original ones.")
-            classes = sorted(set(self._class_mappings.values()))
-        else:
-            classes = all_classes
-        return ["background"] + classes
-
-    @property
-    @override
-    def class_to_idx(self) -> Dict[str, int]:
-        return {label: index for index, label in enumerate(self.classes)}
-
-    @property
-    def _file_suffix(self) -> str:
-        return "nii" if self._decompress else "nii.gz"
-
-    @functools.cached_property
-    def _classes_hash(self) -> str:
-        return hashlib.md5(str(self.classes).encode(), usedforsecurity=False).hexdigest()
-
-    @override
-    def filename(self, index: int) -> str:
-        sample_idx, _ = self._indices[index]
-        sample_dir = self._samples_dirs[sample_idx]
-        return os.path.join(sample_dir, f"ct.{self._file_suffix}")
-
-    @override
-    def prepare_data(self) -> None:
-        if self._download:
-            self._download_dataset()
-        if self._decompress:
-            self._decompress_files()
-        self._samples_dirs = self._fetch_samples_dirs()
-        if self._optimize_mask_loading:
-            self._export_semantic_label_masks()
-
-    @override
-    def configure(self) -> None:
-        self._indices = self._create_indices()
-
-    @override
-    def validate(self) -> None:
-        if self._version is None or self._sample_every_n_slices is not None:
-            return
-
-        if self._classes:
-            last_label = self._classes[-1]
-            n_classes = len(self._classes)
-        elif self._class_mappings:
-            classes = sorted(set(self._class_mappings.values()))
-            last_label = classes[-1]
-            n_classes = len(classes)
-        else:
-            last_label = "vertebrae_T9"
-            n_classes = 117
-
-        _validators.check_dataset_integrity(
-            self,
-            length=self._expected_dataset_lengths.get(f"{self._split}_{self._version}", 0),
-            n_classes=n_classes + 1,
-            first_and_last_labels=("background", last_label),
-        )
-
-    @override
-    def __len__(self) -> int:
-        return len(self._indices)
-
-    @override
-    def load_data(self, index: int) -> tv_tensors.Image:
-        sample_index, slice_index = self._indices[index]
-        image_path = self._get_image_path(sample_index)
-        image_nii = io.read_nifti(image_path, slice_index)
-        image_array = io.nifti_to_array(image_nii)
-        image_array = self._fix_orientation(image_array)
-        return tv_tensors.Image(image_array.copy().transpose(2, 0, 1))
-
-    @override
-    def load_target(self, index: int) -> tv_tensors.Mask:
-        if self._optimize_mask_loading:
-            mask = self._load_semantic_label_mask(index)
-        else:
-            mask = self._load_target(index)
-        mask = self._fix_orientation(mask)
-        return tv_tensors.Mask(mask.copy().squeeze(), dtype=torch.int64)  # type: ignore
-
-    @override
-    def load_metadata(self, index: int) -> Dict[str, Any]:
-        _, slice_index = self._indices[index]
-        return {"slice_index": slice_index}
-
-    def _load_target(self, index: int) -> npt.NDArray[Any]:
-        sample_index, slice_index = self._indices[index]
-        return self._load_masks_as_semantic_label(sample_index, slice_index)
-
-    def _load_semantic_label_mask(self, index: int) -> npt.NDArray[Any]:
-        """Loads the segmentation mask from a semantic label NifTi file."""
-        sample_index, slice_index = self._indices[index]
-        nii = io.read_nifti(self._get_optimized_masks_file(sample_index), slice_index)
-        return io.nifti_to_array(nii)
-
-    def _load_masks_as_semantic_label(
-        self, sample_index: int, slice_index: int | None = None
-    ) -> npt.NDArray[Any]:
-        """Loads binary masks as a semantic label mask.
-
-        Args:
-            sample_index: The data sample index.
-            slice_index: Whether to return only a specific slice.
-        """
-        masks_dir = self._get_masks_dir(sample_index)
-        classes = self._class_mappings.keys() if self._class_mappings else self.classes[1:]
-        mask_paths = [os.path.join(masks_dir, f"{label}.nii.gz") for label in classes]
-        binary_masks = [io.nifti_to_array(io.read_nifti(path, slice_index)) for path in mask_paths]
-
-        if self._class_mappings:
-            mapped_binary_masks = [np.zeros_like(binary_masks[0], dtype=np.bool_)] * len(
-                self.classes[1:]
-            )
-            for original_class, mapped_class in self._class_mappings.items():
-                mapped_index = self.class_to_idx[mapped_class] - 1
-                original_index = list(self._class_mappings.keys()).index(original_class)
-                mapped_binary_masks[mapped_index] = np.logical_or(
-                    mapped_binary_masks[mapped_index], binary_masks[original_index]
-                )
-            binary_masks = mapped_binary_masks
-
-        background_mask = np.zeros_like(binary_masks[0])
-        return np.argmax([background_mask] + binary_masks, axis=0)
-
-    def _export_semantic_label_masks(self) -> None:
-        """Exports the segmentation binary masks (one-hot) to semantic labels."""
-        mask_classes_file = os.path.join(f"{self._get_optimized_masks_root()}/classes.txt")
-        if os.path.isfile(mask_classes_file):
-            with open(mask_classes_file, "r") as file:
-                if file.read() != str(self.classes):
-                    raise ValueError(
-                        "Optimized masks hash doesn't match the current classes or mappings."
-                    )
-            return
-
-        total_samples = len(self._samples_dirs)
-        semantic_labels = [
-            (index, self._get_optimized_masks_file(index)) for index in range(total_samples)
-        ]
-        to_export = filter(lambda x: not os.path.isfile(x[1]), semantic_labels)
-
-        def _process_mask(sample_index: Any, filename: str) -> None:
-            semantic_labels = self._load_masks_as_semantic_label(sample_index)
-            os.makedirs(os.path.dirname(filename), exist_ok=True)
-            io.save_array_as_nifti(semantic_labels, filename)
-
-        multiprocessing.run_with_threads(
-            _process_mask,
-            list(to_export),
-            num_workers=self._num_workers,
-            progress_desc=">> Exporting optimized semantic mask",
-            return_results=False,
-        )
-
-        os.makedirs(os.path.dirname(mask_classes_file), exist_ok=True)
-        with open(mask_classes_file, "w") as file:
-            file.write(str(self.classes))
-
-    def _fix_orientation(self, array: npt.NDArray):
-        """Fixes orientation such that table is at the bottom & liver on the left."""
-        array = np.rot90(array)
-        array = np.flip(array, axis=1)
-        return array
-
-    def _get_image_path(self, sample_index: int) -> str:
-        """Returns the corresponding image path."""
-        sample_dir = self._samples_dirs[sample_index]
-        return os.path.join(self._root, sample_dir, f"ct.{self._file_suffix}")
-
-    def _get_masks_dir(self, sample_index: int) -> str:
-        """Returns the directory of the corresponding masks."""
-        sample_dir = self._samples_dirs[sample_index]
-        return os.path.join(self._root, sample_dir, "segmentations")
-
-    def _get_optimized_masks_root(self) -> str:
-        """Returns the directory of the optimized masks."""
-        return os.path.join(self._root, f"processed/masks/{self._classes_hash}")
-
-    def _get_optimized_masks_file(self, sample_index: int) -> str:
-        """Returns the semantic label filename."""
-        return os.path.join(
-            f"{self._get_optimized_masks_root()}/{self._samples_dirs[sample_index]}/masks.nii"
-        )
-
-    def _get_number_of_slices_per_sample(self, sample_index: int) -> int:
-        """Returns the total amount of slices of a sample."""
-        image_path = self._get_image_path(sample_index)
-        image_shape = io.fetch_nifti_shape(image_path)
-        return image_shape[-1]
-
-    def _fetch_samples_dirs(self) -> List[str]:
-        """Returns the name of all the samples of all the splits of the dataset."""
-        sample_filenames = [
-            filename
-            for filename in os.listdir(self._root)
-            if os.path.isdir(os.path.join(self._root, filename)) and re.match(r"^s\d{4}$", filename)
-        ]
-        return sorted(sample_filenames)
-
-    def _get_split_indices(self) -> List[int]:
-        """Returns the samples indices that corresponding the dataset split and version."""
-        metadata_file = os.path.join(self._root, "meta.csv")
-        metadata = io.read_csv(metadata_file, delimiter=";", encoding="utf-8-sig")
-
-        match self._split:
-            case "train":
-                image_ids = [item["image_id"] for item in metadata if item["split"] == "train"]
-            case "val":
-                image_ids = [item["image_id"] for item in metadata if item["split"] == "val"]
-            case "test":
-                image_ids = [item["image_id"] for item in metadata if item["split"] == "test"]
-            case _:
-                image_ids = self._samples_dirs
-
-        return sorted(map(self._samples_dirs.index, image_ids))
-
-    def _create_indices(self) -> List[Tuple[int, int]]:
-        """Builds the dataset indices for the specified split.
-
-        Returns:
-            A list of tuples, where the first value indicates the
-            sample index which the second its corresponding slice
-            index.
-        """
-        indices = [
-            (sample_idx, slide_idx)
-            for sample_idx in self._get_split_indices()
-            for slide_idx in range(self._get_number_of_slices_per_sample(sample_idx))
-            if slide_idx % (self._sample_every_n_slices or 1) == 0
-        ]
-        return indices
-
-    def _download_dataset(self) -> None:
-        """Downloads the dataset."""
-        dataset_resources = {
-            "small": self._resources_small,
-            "full": self._resources_full,
-        }
-        resources = dataset_resources.get(self._version or "")
-        if resources is None:
-            raise ValueError(
-                f"Can't download data version '{self._version}'. Use 'small' or 'full'."
-            )
-
-        self._print_license()
-        for resource in resources:
-            if os.path.isdir(self._root):
-                continue
-
-            utils.download_and_extract_archive(
-                resource.url,
-                download_root=self._root,
-                filename=resource.filename,
-                remove_finished=True,
-            )
-
-    def _decompress_files(self) -> None:
-        compressed_paths = Path(self._root).rglob("*/ct.nii.gz")
-        multiprocessing.run_with_threads(
-            core_io.gunzip_file,
-            [(str(path),) for path in compressed_paths],
-            num_workers=self._num_workers,
-            progress_desc=">> Decompressing .gz files",
-            return_results=False,
-        )
-
-    def _print_license(self) -> None:
-        """Prints the dataset license."""
-        print(f"Dataset license: {self._license}")