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

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

@@ -0,0 +1,243 @@
+"""Dataset for Task 7 (pancreas tumor) from the Medical Segmentation Decathlon (MSD)."""
+
+import glob
+import os
+import re
+from typing import Any, Callable, Dict, List, Literal, Tuple
+
+import huggingface_hub
+from torchvision import tv_tensors
+from torchvision.datasets import utils as data_utils
+from typing_extensions import override
+
+from eva.core.utils import requirements
+from eva.vision.data import tv_tensors as eva_tv_tensors
+from eva.vision.data.datasets.segmentation import _utils
+from eva.vision.data.datasets.segmentation.metadata import _msd_task7_pancreas
+from eva.vision.data.datasets.vision import VisionDataset
+
+
+class MSDTask7Pancreas(VisionDataset[eva_tv_tensors.Volume, tv_tensors.Mask]):
+    """Task 7 (pancreas tumor) of the Medical Segmentation Decathlon (MSD).
+
+    The data set consists of 420 portal-venous phase CT scans of patients undergoing
+    resection of pancreatic masses. The corresponding target ROIs were the pancreatic
+    parenchyma and pancreatic mass (cyst or tumor). This data set was selected due to
+    label unbalance between large (background), medium (pancreas) and small (tumor)
+    structures. The data was acquired in the Memorial Sloan Kettering Cancer
+    Center, New York, US.
+
+    More info:
+      - Paper: https://www.nature.com/articles/s41467-022-30695-9
+      - Dataset source: https://huggingface.co/datasets/Luffy503/VoCo_Downstream
+    """
+
+    _train_ids = _msd_task7_pancreas.train_ids
+    """File indices of the training set."""
+
+    _val_ids = _msd_task7_pancreas.val_ids
+    """File indices of the validation set."""
+
+    def __init__(
+        self,
+        root: str,
+        split: Literal["train", "val"] | None = None,
+        download: bool = False,
+        transforms: Callable | None = None,
+    ) -> None:
+        """Initializes the dataset.
+
+        Args:
+            root: Path to the dataset root directory.
+            split: Dataset split to use ('train' or 'val').
+                If None, it uses the full dataset.
+            download: Whether to download the dataset.
+            transforms: A callable object for applying data transformations.
+                If None, no transformations are applied.
+        """
+        super().__init__()
+
+        self._root = root
+        self._split = split
+        self._download = download
+        self._transforms = transforms
+
+        self._samples: Dict[int, Tuple[str, str]]
+        self._indices: List[int]
+
+    @property
+    @override
+    def classes(self) -> List[str]:
+        return [
+            "background",
+            "pancreas",
+            "cancer",
+        ]
+
+    @property
+    @override
+    def class_to_idx(self) -> Dict[str, int]:
+        return {label: index for index, label in enumerate(self.classes)}
+
+    @override
+    def filename(self, index: int) -> str:
+        return os.path.relpath(self._samples[self._indices[index]][0], self._root)
+
+    @override
+    def prepare_data(self) -> None:
+        if self._download:
+            self._download_dataset()
+
+    @override
+    def configure(self) -> None:
+        self._samples = self._find_samples()
+        self._indices = self._make_indices()
+
+    @override
+    def validate(self) -> None:
+        requirements.check_dependencies(requirements={"torch": "2.5.1", "torchvision": "0.20.1"})
+
+        def _valid_sample(index: int) -> bool:
+            """Indicates if the sample files exist and are reachable."""
+            volume_file, segmentation_file = self._samples[self._indices[index]]
+            return os.path.isfile(volume_file) and os.path.isfile(segmentation_file)
+
+        if len(self._samples) < len(self._indices):
+            raise OSError(f"Dataset is missing {len(self._indices) - len(self._samples)} files.")
+
+        invalid_samples = [self._samples[i] for i in range(len(self)) if not _valid_sample(i)]
+        if invalid_samples:
+            raise OSError(
+                f"Dataset '{self.__class__.__qualname__}' contains missing or "
+                f"corrupted samples  ({len(invalid_samples)} in total). "
+                f"Examples of missing folders: {str(invalid_samples[:10])[:-1]}, ...]. "
+            )
+
+    @override
+    def __getitem__(
+        self, index: int
+    ) -> tuple[eva_tv_tensors.Volume, tv_tensors.Mask, dict[str, Any]]:
+        volume = self.load_data(index)
+        mask = self.load_target(index)
+        metadata = self.load_metadata(index) or {}
+        volume_tensor, mask_tensor = self._apply_transforms(volume, mask)
+        return volume_tensor, mask_tensor, metadata
+
+    @override
+    def __len__(self) -> int:
+        return len(self._indices)
+
+    @override
+    def load_data(self, index: int) -> eva_tv_tensors.Volume:
+        """Loads the CT volume for a given sample.
+
+        Args:
+            index: The index of the desired sample.
+
+        Returns:
+            Tensor representing the CT volume of shape `[T, C, H, W]`.
+        """
+        ct_scan_file, _ = self._samples[self._indices[index]]
+        return _utils.load_volume_tensor(ct_scan_file)
+
+    @override
+    def load_target(self, index: int) -> tv_tensors.Mask:
+        """Loads the segmentation mask for a given sample.
+
+        Args:
+            index: The index of the desired sample.
+
+        Returns:
+            Tensor representing the segmentation mask of shape `[T, C, H, W]`.
+        """
+        ct_scan_file, mask_file = self._samples[self._indices[index]]
+        return _utils.load_mask_tensor(mask_file, ct_scan_file)
+
+    def _apply_transforms(
+        self, ct_scan: eva_tv_tensors.Volume, mask: tv_tensors.Mask
+    ) -> Tuple[eva_tv_tensors.Volume, tv_tensors.Mask]:
+        """Applies transformations to the provided data.
+
+        Args:
+            ct_scan: The CT volume tensor.
+            mask: The segmentation mask tensor.
+
+        Returns:
+            A tuple containing the transformed CT and mask tensors.
+        """
+        return self._transforms(ct_scan, mask) if self._transforms else (ct_scan, mask)
+
+    def _find_samples(self) -> Dict[int, Tuple[str, str]]:
+        """Retrieves the file paths for the CT volumes and segmentation.
+
+        Returns:
+            The a dictionary mapping the file id to the volume and segmentation paths.
+        """
+
+        def filename_id_volume(filename: str) -> int:
+            matches = re.match(r".*(\d{3})_\d{4}.*", filename)
+            if matches is None:
+                raise ValueError(f"Filename '{filename}' is not valid.")
+            return int(matches.group(1))
+
+        def filename_id_segmentation(filename: str) -> int:
+            matches = re.match(r".*(\d{3}).*", filename)
+            if matches is None:
+                raise ValueError(f"Filename '{filename}' is not valid.")
+            return int(matches.group(1))
+
+        optional_subdir = os.path.join(self._root, "Dataset007_Pancreas")
+        search_dir = optional_subdir if os.path.isdir(optional_subdir) else self._root
+
+        volume_files_pattern = os.path.join(search_dir, "imagesTr", "*.nii.gz")
+        volume_filenames = glob.glob(volume_files_pattern)
+        volume_ids = {filename_id_volume(filename): filename for filename in volume_filenames}
+
+        segmentation_files_pattern = os.path.join(search_dir, "labelsTr", "*.nii.gz")
+        segmentation_filenames = glob.glob(segmentation_files_pattern)
+        segmentation_ids = {
+            filename_id_segmentation(filename): filename for filename in segmentation_filenames
+        }
+
+        return {
+            file_id: (volume_ids[file_id], segmentation_ids[file_id])
+            for file_id in volume_ids.keys()
+        }
+
+    def _make_indices(self) -> List[int]:
+        """Builds the dataset indices for the specified split."""
+        file_ids = []
+        match self._split:
+            case "train":
+                file_ids = self._train_ids
+            case "val":
+                file_ids = self._val_ids
+            case None:
+                file_ids = self._train_ids + self._val_ids
+            case _:
+                raise ValueError("Invalid data split. Use 'train', 'val' or `None`.")
+
+        return file_ids
+
+    def _download_dataset(self) -> None:
+        hf_token = os.getenv("HF_TOKEN")
+        if not hf_token:
+            raise ValueError("Huggingface token required, please set the HF_TOKEN env variable.")
+
+        huggingface_hub.snapshot_download(
+            "Luffy503/VoCo_Downstream",
+            repo_type="dataset",
+            token=hf_token,
+            local_dir=self._root,
+            ignore_patterns=[".git*"],
+            allow_patterns=["**/Dataset007_Pancreas.zip"],
+        )
+
+        zip_path = os.path.join(self._root, "MSD_Decathlon/Dataset007_Pancreas.zip")
+        if not os.path.exists(zip_path):
+            raise FileNotFoundError(
+                f"MSD_Decathlon/Dataset007_Pancreas.zip not found in {self._root}, "
+                "something with the download went wrong."
+            )
+
+        data_utils.extract_archive(zip_path, self._root, remove_finished=True)

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

@@ -18,7 +18,7 @@ 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 import _total_segmentator
+from eva.vision.data.datasets.segmentation.metadata import _total_segmentator
 from eva.vision.utils import io
 
 

eva/vision/data/transforms/__init__.py

@@ -1,7 +1,13 @@
 """Vision data transforms."""
 
-from eva.vision.data.transforms.common import ResizeAndCrop
-from eva.vision.data.transforms.croppad import CropForeground, RandCropByPosNegLabel, SpatialPad
+from eva.vision.data.transforms.common import ResizeAndCrop, Squeeze
+from eva.vision.data.transforms.croppad import (
+    CropForeground,
+    RandCropByLabelClasses,
+    RandCropByPosNegLabel,
+    RandSpatialCrop,
+    SpatialPad,
+)
 from eva.vision.data.transforms.intensity import (
     RandScaleIntensity,
     RandShiftIntensity,
@@ -12,7 +18,9 @@ from eva.vision.data.transforms.utility import EnsureChannelFirst
 
 __all__ = [
     "ResizeAndCrop",
+    "Squeeze",
     "CropForeground",
+    "RandCropByLabelClasses",
     "RandCropByPosNegLabel",
     "SpatialPad",
     "RandScaleIntensity",
@@ -20,6 +28,7 @@ __all__ = [
     "ScaleIntensityRange",
     "RandFlip",
     "RandRotate90",
+    "RandSpatialCrop",
     "Spacing",
     "EnsureChannelFirst",
 ]

eva/vision/data/transforms/base/__init__.py

@@ -0,0 +1,5 @@
+"""Base classes for transforms."""
+
+from eva.vision.data.transforms.base.monai import RandomMonaiTransform
+
+__all__ = ["RandomMonaiTransform"]

eva/vision/data/transforms/base/monai.py

@@ -0,0 +1,27 @@
+"""Base class for MONAI transform wrappers."""
+
+import abc
+
+from torchvision.transforms import v2
+
+
+class RandomMonaiTransform(v2.Transform, abc.ABC):
+    """Base class for MONAI transform wrappers."""
+
+    @abc.abstractmethod
+    def set_random_state(self, seed: int) -> None:
+        """Set the random state for the transform.
+
+        MONAI's random transforms use numpy.random for random number generation
+        which is seeded at the very beginning by lightning's seed_everything, but when
+        torch spins up the dataloader workers, it will only reseed torch's random states
+        and not numpy - so you basically end up with multiple dataloader workers that
+        have equally seeded random transforms, resulting in redundant transform outputs
+        and therefore reducing the diversity of the resulting training data.
+        To solve this, this method should be called in the dataloader's worker_init_fn
+        with a unique seed for each worker.
+
+        Args:
+            seed: The seed to set for the random state of the transform.
+        """
+        pass

eva/vision/data/transforms/common/__init__.py

@@ -1,5 +1,6 @@
 """Common vision transforms."""
 
 from eva.vision.data.transforms.common.resize_and_crop import ResizeAndCrop
+from eva.vision.data.transforms.common.squeeze import Squeeze
 
-__all__ = ["ResizeAndCrop"]
+__all__ = ["ResizeAndCrop", "Squeeze"]

eva/vision/data/transforms/common/squeeze.py

@@ -0,0 +1,24 @@
+"""Squeeze transform."""
+
+from typing import Any
+
+import torch
+from torchvision import tv_tensors
+from torchvision.transforms import v2
+
+
+class Squeeze(v2.Transform):
+    """Squeezes the input tensor accross all or specified dimensions."""
+
+    def __init__(self, dim: int | list[int] | None = None):
+        """Initializes the transform.
+
+        Args:
+            dim: If specified, the input will be squeezed only in the specified dimensions.
+        """
+        super().__init__()
+        self._dim = dim
+
+    def _transform(self, inpt: Any, params: dict[str, Any]) -> Any:
+        output = torch.squeeze(inpt) if self._dim is None else torch.squeeze(inpt, dim=self._dim)
+        return tv_tensors.wrap(output, like=inpt)

eva/vision/data/transforms/croppad/__init__.py

@@ -1,11 +1,15 @@
 """Transforms for crop and pad operations."""
 
 from eva.vision.data.transforms.croppad.crop_foreground import CropForeground
+from eva.vision.data.transforms.croppad.rand_crop_by_label_classes import RandCropByLabelClasses
 from eva.vision.data.transforms.croppad.rand_crop_by_pos_neg_label import RandCropByPosNegLabel
+from eva.vision.data.transforms.croppad.rand_spatial_crop import RandSpatialCrop
 from eva.vision.data.transforms.croppad.spatial_pad import SpatialPad
 
 __all__ = [
     "CropForeground",
+    "RandCropByLabelClasses",
     "RandCropByPosNegLabel",
+    "RandSpatialCrop",
     "SpatialPad",
 ]

eva/vision/data/transforms/croppad/rand_crop_by_label_classes.py

@@ -0,0 +1,74 @@
+"""Crop by label classes transform."""
+
+import functools
+from typing import Any, Dict, List, Sequence
+
+import torch
+from monai.config.type_definitions import NdarrayOrTensor
+from monai.transforms.croppad import array as monai_croppad_transforms
+from torchvision import tv_tensors
+from typing_extensions import override
+
+from eva.vision.data import tv_tensors as eva_tv_tensors
+from eva.vision.data.transforms import base
+
+
+class RandCropByLabelClasses(base.RandomMonaiTransform):
+    """Crop random fixed sized regions with the center belonging to one of the classes.
+
+    Please refer to `monai.transforms.croppad.RandCropByLabelClasses` docs for more details.
+    """
+
+    def __init__(
+        self,
+        spatial_size: Sequence[int] | int,
+        ratios: list[float | int] | None = None,
+        label: torch.Tensor | None = None,
+        num_classes: int | None = None,
+        num_samples: int = 1,
+        image: torch.Tensor | None = None,
+        image_threshold: float = 0.0,
+        indices: list[NdarrayOrTensor] | None = None,
+        allow_smaller: bool = False,
+        warn: bool = True,
+        max_samples_per_class: int | None = None,
+        lazy: bool = False,
+    ) -> None:
+        """Initializes the transform."""
+        super().__init__()
+
+        self._rand_crop = monai_croppad_transforms.RandCropByLabelClasses(
+            spatial_size=spatial_size,
+            ratios=ratios,
+            label=label,
+            num_classes=num_classes,
+            num_samples=num_samples,
+            image=image,
+            image_threshold=image_threshold,
+            indices=indices,
+            allow_smaller=allow_smaller,
+            warn=warn,
+            max_samples_per_class=max_samples_per_class,
+            lazy=lazy,
+        )
+
+    @override
+    def set_random_state(self, seed: int) -> None:
+        self._rand_crop.set_random_state(seed)
+
+    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+        mask = next(inpt for inpt in flat_inputs if isinstance(inpt, tv_tensors.Mask))
+        self._rand_crop.randomize(label=mask)
+        return {}
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_transform.register(tv_tensors.Image)
+    @_transform.register(eva_tv_tensors.Volume)
+    @_transform.register(tv_tensors.Mask)
+    def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        inpt_foreground_crops = self._rand_crop(img=inpt, randomize=False)
+        return [tv_tensors.wrap(crop, like=inpt) for crop in inpt_foreground_crops]

eva/vision/data/transforms/croppad/rand_crop_by_pos_neg_label.py

@@ -7,13 +7,13 @@ import torch
 from monai.config.type_definitions import NdarrayOrTensor
 from monai.transforms.croppad import array as monai_croppad_transforms
 from torchvision import tv_tensors
-from torchvision.transforms import v2
 from typing_extensions import override
 
 from eva.vision.data import tv_tensors as eva_tv_tensors
+from eva.vision.data.transforms import base
 
 
-class RandCropByPosNegLabel(v2.Transform):
+class RandCropByPosNegLabel(base.RandomMonaiTransform):
     """Crop random fixed sized regions with the center being a foreground or background voxel.
 
     Its based on the Pos Neg Ratio and will return a list of arrays for all the cropped images.
@@ -91,6 +91,10 @@ class RandCropByPosNegLabel(v2.Transform):
             lazy=False,
         )
 
+    @override
+    def set_random_state(self, seed: int) -> None:
+        self._rand_crop.set_random_state(seed)
+
     def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
         mask = next(inpt for inpt in flat_inputs if isinstance(inpt, tv_tensors.Mask))
         self._rand_crop.randomize(label=mask)

eva/vision/data/transforms/croppad/rand_spatial_crop.py

@@ -0,0 +1,89 @@
+"""Crop image with random size or specific size ROI."""
+
+import functools
+from typing import Any, Dict, List, Sequence, Tuple
+
+from monai.transforms.croppad import array as monai_croppad_transforms
+from torchvision import tv_tensors
+from torchvision.transforms import v2
+from torchvision.transforms.v2 import _utils as tv_utils
+from typing_extensions import override
+
+from eva.vision.data import tv_tensors as eva_tv_tensors
+
+
+class RandSpatialCrop(v2.Transform):
+    """Crop image with random size or specific size ROI.
+
+    It can crop at a random position as center or at the image center.
+    And allows to set the minimum and maximum size to limit the randomly
+    generated ROI.
+    """
+
+    def __init__(
+        self,
+        roi_size: Sequence[int] | int,
+        max_roi_size: Sequence[int] | int | None = None,
+        random_center: bool = True,
+        random_size: bool = False,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            roi_size: if `random_size` is True, it specifies the minimum crop
+                region. if `random_size` is False, it specifies the expected
+                ROI size to crop. e.g. [224, 224, 128]. if a dimension of ROI
+                size is larger than image size, will not crop that dimension of
+                the image. If its components have non-positive values, the
+                corresponding size of input image will be used. for example: if
+                the spatial size of input data is [40, 40, 40] and
+                `roi_size=[32, 64, -1]`, the spatial size of output data will be
+                [32, 40, 40].
+            max_roi_size: if `random_size` is True and `roi_size` specifies the
+                min crop region size, `max_roi_size` can specify the max crop
+                region size. if None, defaults to the input image size. if its
+                components have non-positive values, the corresponding size of
+                input image will be used.
+            random_center: crop at random position as center or the image center.
+            random_size: crop with random size or specific size ROI. if True, the
+                actual size is sampled from `randint(roi_size, max_roi_size + 1)`.
+        """
+        super().__init__()
+
+        self._rand_spatial_crop = monai_croppad_transforms.RandSpatialCrop(
+            roi_size=roi_size,
+            max_roi_size=max_roi_size,
+            random_center=random_center,
+            random_size=random_size,
+        )
+        self._cropper = monai_croppad_transforms.Crop()
+
+    def set_random_state(self, seed: int) -> None:
+        """Set the random state for the transform."""
+        self._rand_spatial_crop.set_random_state(seed)
+
+    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+        t, h, w = tv_utils.query_chw(flat_inputs)
+        self._rand_spatial_crop.randomize((t, h, w))
+        return {}
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_transform.register(tv_tensors.Image)
+    @_transform.register(eva_tv_tensors.Volume)
+    @_transform.register(tv_tensors.Mask)
+    def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        slices = self._get_crop_slices()
+        inpt_rand_crop = self._cropper(inpt, slices=slices)
+        return tv_tensors.wrap(inpt_rand_crop, like=inpt)
+
+    def _get_crop_slices(self) -> Tuple[slice, ...]:
+        """Returns the sequence of slices to crop."""
+        if self._rand_spatial_crop.random_center:
+            return self._rand_spatial_crop._slices
+
+        central_cropper = monai_croppad_transforms.CenterSpatialCrop(self._size)
+        return central_cropper.compute_slices(self._rand_spatial_crop._size)  # type: ignore

eva/vision/data/transforms/intensity/rand_scale_intensity.py

@@ -7,13 +7,13 @@ import numpy as np
 from monai.config.type_definitions import DtypeLike
 from monai.transforms.intensity import array as monai_intensity_transforms
 from torchvision import tv_tensors
-from torchvision.transforms import v2
 from typing_extensions import override
 
 from eva.vision.data import tv_tensors as eva_tv_tensors
+from eva.vision.data.transforms import base
 
 
-class RandScaleIntensity(v2.Transform):
+class RandScaleIntensity(base.RandomMonaiTransform):
     """Randomly scale the intensity of input image.
 
     The factor is by ``v = v * (1 + factor)``, where
@@ -47,6 +47,10 @@ class RandScaleIntensity(v2.Transform):
             dtype=dtype,
         )
 
+    @override
+    def set_random_state(self, seed: int) -> None:
+        self._rand_scale_intensity.set_random_state(seed)
+
     @functools.singledispatchmethod
     @override
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:

eva/vision/data/transforms/intensity/rand_shift_intensity.py

@@ -5,13 +5,13 @@ from typing import Any, Dict
 
 from monai.transforms.intensity import array as monai_intensity_transforms
 from torchvision import tv_tensors
-from torchvision.transforms import v2
 from typing_extensions import override
 
 from eva.vision.data import tv_tensors as eva_tv_tensors
+from eva.vision.data.transforms import base
 
 
-class RandShiftIntensity(v2.Transform):
+class RandShiftIntensity(base.RandomMonaiTransform):
     """Randomly shift intensity with randomly picked offset."""
 
     def __init__(
@@ -36,13 +36,17 @@ class RandShiftIntensity(v2.Transform):
         """
         super().__init__()
 
-        self._rand_swift_intensity = monai_intensity_transforms.RandShiftIntensity(
+        self._rand_shift_intensity = monai_intensity_transforms.RandShiftIntensity(
             offsets=offsets,
             safe=safe,
             prob=prob,
             channel_wise=channel_wise,
         )
 
+    @override
+    def set_random_state(self, seed: int) -> None:
+        self._rand_shift_intensity.set_random_state(seed)
+
     @functools.singledispatchmethod
     @override
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
@@ -51,5 +55,5 @@ class RandShiftIntensity(v2.Transform):
     @_transform.register(tv_tensors.Image)
     @_transform.register(eva_tv_tensors.Volume)
     def _(self, inpt: tv_tensors.Image, params: Dict[str, Any]) -> Any:
-        inpt_scaled = self._rand_swift_intensity(inpt)
+        inpt_scaled = self._rand_shift_intensity(inpt)
         return tv_tensors.wrap(inpt_scaled, like=inpt)