kaiko-eva 0.0.0.dev6__py3-none-any.whl

eva/vision/data/datasets/classification/patch_camelyon.py

@@ -0,0 +1,203 @@
+"""PatchCamelyon dataset."""
+
+import os
+from typing import Callable, Dict, List, Literal
+
+import h5py
+import numpy as np
+from torchvision.datasets import utils
+from typing_extensions import override
+
+from eva.vision.data.datasets import _validators, structs
+from eva.vision.data.datasets.classification import base
+
+_URL_TEMPLATE = "https://zenodo.org/records/2546921/files/{filename}.gz?download=1"
+"""PatchCamelyon URL files templates."""
+
+
+class PatchCamelyon(base.ImageClassification):
+    """Dataset class for PatchCamelyon images and corresponding targets."""
+
+    _train_resources: List[structs.DownloadResource] = [
+        structs.DownloadResource(
+            filename="camelyonpatch_level_2_split_train_x.h5",
+            url=_URL_TEMPLATE.format(filename="camelyonpatch_level_2_split_train_x.h5"),
+            md5="01844da899645b4d6f84946d417ba453",
+        ),
+        structs.DownloadResource(
+            filename="camelyonpatch_level_2_split_train_y.h5",
+            url=_URL_TEMPLATE.format(filename="camelyonpatch_level_2_split_train_y.h5"),
+            md5="0781386bf6c2fb62d58ff18891466aca",
+        ),
+    ]
+    """Train resources."""
+
+    _val_resources: List[structs.DownloadResource] = [
+        structs.DownloadResource(
+            filename="camelyonpatch_level_2_split_valid_x.h5",
+            url=_URL_TEMPLATE.format(filename="camelyonpatch_level_2_split_valid_x.h5"),
+            md5="81cf9680f1724c40673f10dc88e909b1",
+        ),
+        structs.DownloadResource(
+            filename="camelyonpatch_level_2_split_valid_y.h5",
+            url=_URL_TEMPLATE.format(filename="camelyonpatch_level_2_split_valid_y.h5"),
+            md5="94d8aacc249253159ce2a2e78a86e658",
+        ),
+    ]
+    """Validation resources."""
+
+    _test_resources: List[structs.DownloadResource] = [
+        structs.DownloadResource(
+            filename="camelyonpatch_level_2_split_test_x.h5",
+            url=_URL_TEMPLATE.format(filename="camelyonpatch_level_2_split_test_x.h5"),
+            md5="2614b2e6717d6356be141d9d6dbfcb7e",
+        ),
+        structs.DownloadResource(
+            filename="camelyonpatch_level_2_split_test_y.h5",
+            url=_URL_TEMPLATE.format(filename="camelyonpatch_level_2_split_test_y.h5"),
+            md5="11ed647efe9fe457a4eb45df1dba19ba",
+        ),
+    ]
+    """Test resources."""
+
+    _license: str = (
+        "Creative Commons Zero v1.0 Universal (https://choosealicense.com/licenses/cc0-1.0/)"
+    )
+    """Dataset license."""
+
+    def __init__(
+        self,
+        root: str,
+        split: Literal["train", "val", "test"],
+        download: bool = False,
+        image_transforms: Callable | None = None,
+        target_transforms: Callable | None = None,
+    ) -> None:
+        """Initializes the dataset.
+
+        Args:
+            root: The path to the dataset root. This path should contain
+                the uncompressed h5 files and the metadata.
+            split: The dataset split for training, validation, or testing.
+            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.
+            image_transforms: A function/transform that takes in an image
+                and returns a transformed version.
+            target_transforms: A function/transform that takes in the target
+                and transforms it.
+        """
+        super().__init__(
+            image_transforms=image_transforms,
+            target_transforms=target_transforms,
+        )
+
+        self._root = root
+        self._split = split
+        self._download = download
+
+    @property
+    @override
+    def classes(self) -> List[str]:
+        return ["no_tumor", "tumor"]
+
+    @property
+    @override
+    def class_to_idx(self) -> Dict[str, int]:
+        return {"no_tumor": 0, "tumor": 1}
+
+    @override
+    def filename(self, index: int) -> str:
+        return f"camelyonpatch_level_2_split_{self._split}_x_{index}"
+
+    @override
+    def prepare_data(self) -> None:
+        if self._download:
+            self._download_dataset()
+
+    @override
+    def validate(self) -> None:
+        expected_length = {
+            "train": 262144,
+            "val": 32768,
+            "test": 32768,
+        }
+        _validators.check_dataset_integrity(
+            self,
+            length=expected_length.get(self._split, 0),
+            n_classes=2,
+            first_and_last_labels=("no_tumor", "tumor"),
+        )
+
+    @override
+    def load_image(self, index: int) -> np.ndarray:
+        return self._load_from_h5("x", index)
+
+    @override
+    def load_target(self, index: int) -> np.ndarray:
+        target = self._load_from_h5("y", index).squeeze()
+        return np.asarray(target, dtype=np.int64)
+
+    @override
+    def __len__(self) -> int:
+        return self._fetch_dataset_length()
+
+    def _download_dataset(self) -> None:
+        """Downloads the PatchCamelyon dataset."""
+        for resource in self._train_resources + self._val_resources + self._test_resources:
+            file_path = os.path.join(self._root, resource.filename)
+            if utils.check_integrity(file_path, resource.md5):
+                continue
+
+            self._print_license()
+            utils.download_and_extract_archive(
+                resource.url,
+                download_root=self._root,
+                filename=resource.filename + ".gz",
+                remove_finished=True,
+            )
+
+    def _load_from_h5(
+        self,
+        data_key: Literal["x", "y"],
+        index: int | None = None,
+    ) -> np.ndarray:
+        """Load data or targets from an HDF5 file.
+
+        Args:
+            data_key: Specify whether to load 'x' or 'y'.
+            index: Optional parameter to load data/targets at a specific index.
+                If `None`, the entire data/targets array is returned.
+
+        Returns:
+            A array containing the specified data.
+        """
+        h5_file = self._h5_file(data_key)
+        with h5py.File(h5_file, "r") as file:
+            data = file[data_key]
+            return data[:] if index is None else data[index]  # type: ignore
+
+    def _fetch_dataset_length(self) -> int:
+        """Fetches the dataset split length from its HDF5 file."""
+        h5_file = self._h5_file("y")
+        with h5py.File(h5_file, "r") as file:
+            data = file["y"]
+            return len(data)  # type: ignore
+
+    def _h5_file(self, datatype: Literal["x", "y"]) -> str:
+        """Generates the filename for the H5 file based on the specified data type and split.
+
+        Args:
+            datatype: The type of data, where "x" and "y" represent the input
+                and target datasets respectively.
+
+        Returns:
+            The relative file path for the H5 file based on the provided data type and split.
+        """
+        split_suffix = "valid" if self._split == "val" else self._split
+        filename = f"camelyonpatch_level_2_split_{split_suffix}_{datatype}.h5"
+        return os.path.join(self._root, filename)
+
+    def _print_license(self) -> None:
+        """Prints the dataset license."""
+        print(f"Dataset license: {self._license}")

eva/vision/data/datasets/classification/total_segmentator.py

@@ -0,0 +1,212 @@
+"""TotalSegmentator 2D segmentation dataset class."""
+
+import functools
+import os
+from glob import glob
+from typing import Callable, Dict, List, Literal, Tuple
+
+import numpy as np
+from torchvision.datasets import utils
+from typing_extensions import override
+
+from eva.vision.data.datasets import _utils, _validators, structs
+from eva.vision.data.datasets.classification import base
+from eva.vision.utils import io
+
+
+class TotalSegmentatorClassification(base.ImageClassification):
+    """TotalSegmentator multi-label classification dataset."""
+
+    _train_index_ranges: List[Tuple[int, int]] = [(0, 83)]
+    """Train range indices."""
+
+    _val_index_ranges: List[Tuple[int, int]] = [(83, 103)]
+    """Validation range indices."""
+
+    _n_slices_per_image: int = 20
+    """The amount of slices to 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."""
+
+    def __init__(
+        self,
+        root: str,
+        split: Literal["train", "val"] | None,
+        version: Literal["small", "full"] = "small",
+        download: bool = False,
+        image_transforms: Callable | None = None,
+        target_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.
+            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.
+            image_transforms: A function/transform that takes in an image
+                and returns a transformed version.
+            target_transforms: A function/transform that takes in the target
+                and transforms it.
+        """
+        super().__init__(
+            image_transforms=image_transforms,
+            target_transforms=target_transforms,
+        )
+
+        self._root = root
+        self._split = split
+        self._version = version
+        self._download = download
+
+        self._samples_dirs: List[str] = []
+        self._indices: List[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, self._samples_dirs[0], "segmentations", "*.nii.gz"
+        )
+        return sorted(map(get_filename, glob(first_sample_labels)))
+
+    @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:
+        sample_dir = self._samples_dirs[self._indices[index]]
+        return os.path.join(sample_dir, "ct.nii.gz")
+
+    @override
+    def prepare_data(self) -> None:
+        if self._download:
+            self._download_dataset()
+
+    @override
+    def configure(self) -> None:
+        self._samples_dirs = self._fetch_samples_dirs()
+        self._indices = self._create_indices()
+
+    @override
+    def validate(self) -> None:
+        _validators.check_dataset_integrity(
+            self,
+            length=1660 if self._split == "train" else 400,
+            n_classes=117,
+            first_and_last_labels=("adrenal_gland_left", "vertebrae_T9"),
+        )
+
+    @override
+    def __len__(self) -> int:
+        return len(self._indices) * self._n_slices_per_image
+
+    @override
+    def load_image(self, index: int) -> np.ndarray:
+        image_path = self._get_image_path(index)
+        slice_index = self._get_sample_slice_index(index)
+        image_array = io.read_nifti_slice(image_path, slice_index)
+        return image_array.repeat(3, axis=2)
+
+    @override
+    def load_target(self, index: int) -> np.ndarray:
+        masks = self._load_masks(index)
+        targets = [1 in masks[..., mask_index] for mask_index in range(masks.shape[-1])]
+        return np.asarray(targets, dtype=np.int64)
+
+    def _load_masks(self, index: int) -> np.ndarray:
+        """Returns the `index`'th target mask sample."""
+        masks_dir = self._get_masks_dir(index)
+        slice_index = self._get_sample_slice_index(index)
+        mask_paths = (os.path.join(masks_dir, label + ".nii.gz") for label in self.classes)
+        masks = [io.read_nifti_slice(path, slice_index) for path in mask_paths]
+        return np.concatenate(masks, axis=-1)
+
+    def _get_masks_dir(self, index: int) -> str:
+        """Returns the directory of the corresponding masks."""
+        sample_dir = self._get_sample_dir(index)
+        return os.path.join(self._root, sample_dir, "segmentations")
+
+    def _get_image_path(self, index: int) -> str:
+        """Returns the corresponding image path."""
+        sample_dir = self._get_sample_dir(index)
+        return os.path.join(self._root, sample_dir, "ct.nii.gz")
+
+    def _get_sample_dir(self, index: int) -> str:
+        """Returns the corresponding sample directory."""
+        sample_index = self._indices[index // self._n_slices_per_image]
+        return self._samples_dirs[sample_index]
+
+    def _get_sample_slice_index(self, index: int) -> int:
+        """Returns the corresponding slice index."""
+        image_path = self._get_image_path(index)
+        total_slices = io.fetch_total_nifti_slices(image_path)
+        slice_indices = np.linspace(0, total_slices - 1, num=self._n_slices_per_image, dtype=int)
+        return slice_indices[index % self._n_slices_per_image]
+
+    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))
+        ]
+        return sorted(sample_filenames)
+
+    def _create_indices(self) -> List[int]:
+        """Builds the dataset indices for the specified split."""
+        split_index_ranges = {
+            "train": self._train_index_ranges,
+            "val": self._val_index_ranges,
+            None: [(0, 103)],
+        }
+        index_ranges = split_index_ranges.get(self._split)
+        if index_ranges is None:
+            raise ValueError("Invalid data split. Use 'train', 'val' or `None`.")
+
+        return _utils.ranges_to_indices(index_ranges)
+
+    def _download_dataset(self) -> None:
+        """Downloads the dataset."""
+        dataset_resources = {
+            "small": self._resources_small,
+            "full": self._resources_full,
+            None: (0, 103),
+        }
+        resources = dataset_resources.get(self._version)
+        if resources is None:
+            raise ValueError("Invalid data version. Use 'small' or 'full'.")
+
+        for resource in resources:
+            utils.download_and_extract_archive(
+                resource.url,
+                download_root=self._root,
+                filename=resource.filename,
+                remove_finished=True,
+            )

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

@@ -0,0 +1,6 @@
+"""Segmentation datasets API."""
+
+from eva.vision.data.datasets.segmentation.base import ImageSegmentation
+from eva.vision.data.datasets.segmentation.total_segmentator import TotalSegmentator2D
+
+__all__ = ["ImageSegmentation", "TotalSegmentator2D"]

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

@@ -0,0 +1,112 @@
+"""Base for image segmentation datasets."""
+
+import abc
+from typing import Any, Callable, Dict, List, Tuple
+
+import numpy as np
+from typing_extensions import override
+
+from eva.vision.data.datasets import vision
+
+
+class ImageSegmentation(vision.VisionDataset[Tuple[np.ndarray, np.ndarray]], abc.ABC):
+    """Image segmentation abstract dataset."""
+
+    def __init__(
+        self,
+        image_transforms: Callable | None = None,
+        target_transforms: Callable | None = None,
+        image_target_transforms: Callable | None = None,
+    ) -> None:
+        """Initializes the image segmentation base class.
+
+        Args:
+            image_transforms: A function/transform that takes in an image
+                and returns a transformed version.
+            target_transforms: A function/transform that takes in the target
+                and transforms it.
+            image_target_transforms: A function/transforms that takes in an
+                image and a label and returns the transformed versions of both.
+                This transform happens after the `image_transforms` and
+                `target_transforms`.
+        """
+        super().__init__()
+
+        self._image_transforms = image_transforms
+        self._target_transforms = target_transforms
+        self._image_target_transforms = image_target_transforms
+
+    @property
+    def classes(self) -> List[str] | None:
+        """Returns the list with names of the dataset names."""
+
+    @property
+    def class_to_idx(self) -> Dict[str, int] | None:
+        """Returns a mapping of the class name to its target index."""
+
+    def load_metadata(self, index: int | None) -> Dict[str, Any] | List[Dict[str, Any]] | None:
+        """Returns the dataset metadata.
+
+        Args:
+            index: The index of the data sample to return the metadata of.
+                If `None`, it will return the metadata of the current dataset.
+
+        Returns:
+            The sample metadata.
+        """
+
+    @abc.abstractmethod
+    def load_image(self, index: int) -> np.ndarray:
+        """Loads and returns the `index`'th image sample.
+
+        Args:
+            index: The index of the data sample to load.
+
+        Returns:
+            The image as a numpy array.
+        """
+
+    @abc.abstractmethod
+    def load_mask(self, index: int) -> np.ndarray:
+        """Returns the `index`'th target mask sample.
+
+        Args:
+            index: The index of the data sample target mask to load.
+
+        Returns:
+            The sample mask as a stack of binary mask arrays (label, height, width).
+        """
+
+    @abc.abstractmethod
+    @override
+    def __len__(self) -> int:
+        raise NotImplementedError
+
+    @override
+    def __getitem__(self, index: int) -> Tuple[np.ndarray, np.ndarray]:
+        image = self.load_image(index)
+        mask = self.load_mask(index)
+        return self._apply_transforms(image, mask)
+
+    def _apply_transforms(
+        self, image: np.ndarray, target: np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """Applies the transforms to the provided data and returns them.
+
+        Args:
+            image: The desired image.
+            target: The target of the image.
+
+        Returns:
+            A tuple with the image and the target transformed.
+        """
+        if self._image_transforms is not None:
+            image = self._image_transforms(image)
+
+        if self._target_transforms is not None:
+            target = self._target_transforms(target)
+
+        if self._image_target_transforms is not None:
+            image, target = self._image_target_transforms(image, target)
+
+        return image, target