kaiko-eva 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,236 @@
+"""BCSS dataset."""
+
+import glob
+import os
+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.transforms.v2 import functional
+from typing_extensions import override
+
+from eva.vision.data.datasets import _validators, wsi
+from eva.vision.data.datasets.segmentation import _utils, base
+from eva.vision.data.wsi.patching import samplers
+from eva.vision.utils import io
+
+
+class BCSS(wsi.MultiWsiDataset, base.ImageSegmentation):
+    """Dataset class for BCSS semantic segmentation task.
+
+    Source: https://github.com/PathologyDataScience/BCSS
+
+    We apply the the class grouping proposed by the challenge baseline:
+    https://bcsegmentation.grand-challenge.org/Baseline/
+
+    outside_roi: outside_roi
+    tumor: angioinvasion, dcis
+    stroma: stroma
+    inflammatory: lymphocytic_infiltrate, plasma_cells, other_immune_infiltrate
+    necrosis: necrosis_or_debris
+    other: remaining
+
+    Be aware that outside_roi should be assigned zero-weight during model training.
+    """
+
+    _train_split_ratio: float = 0.8
+    """Train split ratio."""
+
+    _val_split_ratio: float = 0.2
+    """Validation split ratio."""
+
+    _expected_length: int = 151
+    """Expected dataset length."""
+
+    _val_institutes = {"BH", "C8", "A8", "A1", "E9"}
+    """Medical institutes to use for the validation split."""
+
+    _test_institutes = {"OL", "LL", "E2", "EW", "GM", "S3"}
+    """Medical institutes to use for the test split."""
+
+    def __init__(
+        self,
+        root: str,
+        sampler: samplers.Sampler,
+        split: Literal["train", "val", "trainval", "test"] | None = None,
+        width: int = 224,
+        height: int = 224,
+        target_mpp: float = 0.5,
+        transforms: Callable | None = None,
+    ) -> None:
+        """Initializes the dataset.
+
+        Args:
+            root: Root directory of the dataset.
+            sampler: The sampler to use for sampling patch coordinates.
+                If `None`, it will use the ::class::`GridSampler` sampler.
+            split: Dataset split to use. If `None`, the entire dataset is used.
+            width: Width of the patches to be extracted, in pixels.
+            height: Height of the patches to be extracted, in pixels.
+            target_mpp: Target microns per pixel (mpp) for the patches.
+            backend: The backend to use for reading the whole-slide images.
+            transforms: Transforms to apply to the extracted image & mask patches.
+        """
+        self._split = split
+        self._root = root
+
+        self.datasets: List[wsi.WsiDataset]  # type: ignore
+
+        wsi.MultiWsiDataset.__init__(
+            self,
+            root=root,
+            file_paths=self._load_file_paths(split),
+            width=width,
+            height=height,
+            sampler=sampler or samplers.GridSampler(max_samples=1000),
+            target_mpp=target_mpp,
+            overwrite_mpp=0.25,
+            backend="pil",
+        )
+        base.ImageSegmentation.__init__(self, transforms=transforms)
+
+    @property
+    @override
+    def classes(self) -> List[str]:
+        return list(self.class_to_idx.keys())
+
+    @property
+    @override
+    def class_to_idx(self) -> Dict[str, int]:
+        return {
+            "outside_roi": 0,
+            "tumor": 1,
+            "stroma": 2,
+            "inflammatory": 3,
+            "necrosis": 4,
+            "other": 5,
+        }
+
+    @override
+    def prepare_data(self) -> None:
+        _validators.check_dataset_exists(self._root, True)
+
+        if not os.path.isdir(os.path.join(self._root, "masks")):
+            raise FileNotFoundError(f"'masks' directory not found in {self._root}.")
+        if not os.path.isdir(os.path.join(self._root, "rgbs_colorNormalized")):
+            raise FileNotFoundError(f"'rgbs_colorNormalized' directory not found in {self._root}.")
+
+    @override
+    def validate(self) -> None:
+        _validators.check_dataset_integrity(
+            self,
+            length=None,
+            n_classes=6,
+            first_and_last_labels=((self.classes[0], self.classes[-1])),
+        )
+
+    @override
+    def __getitem__(self, index: int) -> Tuple[tv_tensors.Image, tv_tensors.Mask, Dict[str, Any]]:
+        return base.ImageSegmentation.__getitem__(self, index)
+
+    @override
+    def load_image(self, index: int) -> tv_tensors.Image:
+        image_array = wsi.MultiWsiDataset.__getitem__(self, index)
+        return functional.to_image(image_array)
+
+    @override
+    def load_mask(self, index: int) -> tv_tensors.Mask:
+        path = self._get_mask_path(index)
+        mask = io.read_image_as_array(path)
+        mask_patch = _utils.extract_mask_patch(mask, self, index)
+        mask_patch = self._map_classes(mask_patch)
+        return tv_tensors.Mask(mask_patch, dtype=torch.int64)  # type: ignore[reportCallIssue]
+
+    @override
+    def load_metadata(self, index: int) -> Dict[str, Any]:
+        (x, y), width, height = _utils.get_coords_at_index(self, index)
+        return {"coords": f"{x},{y},{width},{height}"}
+
+    def _load_file_paths(
+        self, split: Literal["train", "val", "trainval", "test"] | None = None
+    ) -> List[str]:
+        """Loads the file paths of the corresponding dataset split."""
+        file_paths = sorted(glob.glob(os.path.join(self._root, "rgbs_colorNormalized/*.png")))
+        if len(file_paths) != self._expected_length:
+            raise ValueError(
+                f"Expected {self._expected_length} images, found {len(file_paths)} in {self._root}."
+            )
+
+        train_indices, val_indices, test_indices = [], [], []
+        for i, path in enumerate(file_paths):
+            institute = Path(path).stem.split("-")[1]
+            if institute in self._test_institutes:
+                test_indices.append(i)
+            elif institute in self._val_institutes:
+                val_indices.append(i)
+            else:
+                train_indices.append(i)
+
+        match split:
+            case "train":
+                return [file_paths[i] for i in train_indices]
+            case "val":
+                return [file_paths[i] for i in val_indices]
+            case "trainval":
+                return [file_paths[i] for i in train_indices + val_indices]
+            case "test":
+                return [file_paths[i] for i in test_indices]
+            case None:
+                return file_paths
+            case _:
+                raise ValueError("Invalid split. Use 'train', 'val', 'test' or `None`.")
+
+    def _get_mask_path(self, index):
+        """Returns the path to the mask file corresponding to the patch at the given index."""
+        return os.path.join(self._root, "masks", self.filename(index))
+
+    def _map_classes(self, array: npt.NDArray[Any]) -> npt.NDArray[Any]:
+        """Maps the classes of the mask array to the grouped tissue type classes."""
+        original_to_grouped_class_mapping = {
+            "outside_roi": "outside_roi",
+            "angioinvasion": "tumor",
+            "dcis": "tumor",
+            "stroma": "stroma",
+            "lymphocytic_infiltrate": "inflammatory",
+            "plasma_cells": "inflammatory",
+            "other_immune_infiltrate": "inflammatory",
+            "necrosis_or_debris": "necrosis",
+        }
+
+        mapped_array = np.full_like(array, fill_value=self.class_to_idx["other"], dtype=int)
+
+        for original_class, grouped_class in original_to_grouped_class_mapping.items():
+            original_class_idx = _original_class_to_idx[original_class]
+            grouped_class_idx = self.class_to_idx[grouped_class]
+            mapped_array[array == original_class_idx] = grouped_class_idx
+
+        return mapped_array
+
+
+_original_class_to_idx = {
+    "outside_roi": 0,
+    "tumor": 1,
+    "stroma": 2,
+    "lymphocytic_infiltrate": 3,
+    "necrosis_or_debris": 4,
+    "glandular_secretions": 5,
+    "blood": 6,
+    "exclude": 7,
+    "metaplasia_NOS": 8,
+    "fat": 9,
+    "plasma_cells": 10,
+    "other_immune_infiltrate": 11,
+    "mucoid_material": 12,
+    "normal_acinus_or_duct": 13,
+    "lymphatics": 14,
+    "undetermined": 15,
+    "nerve": 16,
+    "skin_adnexa": 17,
+    "blood_vessel": 18,
+    "angioinvasion": 19,
+    "dcis": 20,
+    "other": 21,
+}

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

@@ -0,0 +1,156 @@
+"""CoNSeP dataset."""
+
+import glob
+import os
+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.transforms.v2 import functional
+from typing_extensions import override
+
+from eva.vision.data.datasets import _validators, wsi
+from eva.vision.data.datasets.segmentation import _utils, base
+from eva.vision.data.wsi.patching import samplers
+from eva.vision.utils import io
+
+
+class CoNSeP(wsi.MultiWsiDataset, base.ImageSegmentation):
+    """Dataset class for CoNSeP semantic segmentation task.
+
+    We combine classes 3 (healthy epithelial) & 4 (dysplastic/malignant epithelial)
+    into the epithelial class and 5 (fibroblast), 6 (muscle) & 7 (endothelial) into
+    the spindle-shaped class.
+    """
+
+    _expected_dataset_lengths: Dict[str | None, int] = {
+        "train": 27,
+        "val": 14,
+        None: 41,
+    }
+    """Expected dataset lengths for the splits and complete dataset."""
+
+    def __init__(
+        self,
+        root: str,
+        sampler: samplers.Sampler | None = None,
+        split: Literal["train", "val"] | None = None,
+        width: int = 224,
+        height: int = 224,
+        target_mpp: float = 0.25,
+        transforms: Callable | None = None,
+    ) -> None:
+        """Initializes the dataset.
+
+        Args:
+            root: Root directory of the dataset.
+            sampler: The sampler to use for sampling patch coordinates.
+                If `None`, it will use the ::class::`ForegroundGridSampler` sampler.
+            split: Dataset split to use. If `None`, the entire dataset is used.
+            width: Width of the patches to be extracted, in pixels.
+            height: Height of the patches to be extracted, in pixels.
+            target_mpp: Target microns per pixel (mpp) for the patches.
+            backend: The backend to use for reading the whole-slide images.
+            transforms: Transforms to apply to the extracted image & mask patches.
+        """
+        self._split = split
+        self._root = root
+
+        self.datasets: List[wsi.WsiDataset]  # type: ignore
+
+        wsi.MultiWsiDataset.__init__(
+            self,
+            root=root,
+            file_paths=self._load_file_paths(split),
+            width=width,
+            height=height,
+            sampler=sampler or samplers.ForegroundGridSampler(max_samples=25),
+            target_mpp=target_mpp,
+            overwrite_mpp=0.25,
+            backend="pil",
+            image_transforms=transforms,
+        )
+
+    @property
+    @override
+    def classes(self) -> List[str]:
+        return [
+            "background",
+            "other",
+            "inflammatory",
+            "epithelial",
+            "spindle-shaped",
+        ]
+
+    @property
+    @override
+    def class_to_idx(self) -> Dict[str, int]:
+        return {label: index for index, label in enumerate(self.classes)}
+
+    @override
+    def prepare_data(self) -> None:
+        _validators.check_dataset_exists(self._root, True)
+
+        if not os.path.isdir(os.path.join(self._root, "Train")):
+            raise FileNotFoundError(f"Train directory not found in {self._root}.")
+        if not os.path.isdir(os.path.join(self._root, "Test")):
+            raise FileNotFoundError(f"Test directory not found in {self._root}.")
+
+    @override
+    def validate(self) -> None:
+        _validators.check_dataset_integrity(
+            self,
+            length=None,
+            n_classes=5,
+            first_and_last_labels=((self.classes[0], self.classes[-1])),
+        )
+
+    @override
+    def __getitem__(self, index: int) -> Tuple[tv_tensors.Image, tv_tensors.Mask, Dict[str, Any]]:
+        return base.ImageSegmentation.__getitem__(self, index)
+
+    @override
+    def load_image(self, index: int) -> tv_tensors.Image:
+        image_array = wsi.MultiWsiDataset.__getitem__(self, index)
+        return functional.to_image(image_array)
+
+    @override
+    def load_mask(self, index: int) -> tv_tensors.Mask:
+        path = self._get_mask_path(index)
+        mask = np.array(io.read_mat(path)["type_map"])
+        mask_patch = _utils.extract_mask_patch(mask, self, index)
+        mask_patch = self._map_classes(mask_patch)
+        mask_tensor = tv_tensors.Mask(mask_patch, dtype=torch.int64)  # type: ignore[reportCallIssue]
+        return self._image_transforms(mask_tensor) if self._image_transforms else mask_tensor
+
+    @override
+    def load_metadata(self, index: int) -> Dict[str, Any]:
+        (x, y), width, height = _utils.get_coords_at_index(self, index)
+        return {"coords": f"{x},{y},{width},{height}"}
+
+    def _load_file_paths(self, split: Literal["train", "val"] | None = None) -> List[str]:
+        """Loads the file paths of the corresponding dataset split."""
+        paths = list(glob.glob(os.path.join(self._root, "**/Images/*.png"), recursive=True))
+        n_expected = self._expected_dataset_lengths[None]
+        if len(paths) != n_expected:
+            raise ValueError(f"Expected {n_expected} images, found {len(paths)} in {self._root}.")
+
+        if split is not None:
+            split_to_folder = {"train": "Train", "val": "Test"}
+            paths = filter(lambda p: split_to_folder[split] == p.split("/")[-3], paths)
+
+        return sorted(paths)
+
+    def _get_mask_path(self, index: int) -> str:
+        """Returns the path to the mask file corresponding to the patch at the given index."""
+        filename = self.filename(index).split(".")[0]
+        mask_dir = "Train" if filename.startswith("train") else "Test"
+        return os.path.join(self._root, mask_dir, "Labels", f"{filename}.mat")
+
+    def _map_classes(self, array: npt.NDArray[Any]) -> npt.NDArray[Any]:
+        """Summarizes classes 3 & 4, and 5, 6."""
+        array = np.where(array == 4, 3, array)
+        array = np.where(array > 4, 4, array)
+        return array

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

@@ -0,0 +1,34 @@
+"""Embeddings based semantic segmentation dataset."""
+
+import os
+from typing import List
+
+import torch
+from torchvision import tv_tensors
+from typing_extensions import override
+
+from eva.core.data.datasets import embeddings as embeddings_base
+
+
+class EmbeddingsSegmentationDataset(embeddings_base.EmbeddingsDataset[tv_tensors.Mask]):
+    """Embeddings segmentation dataset."""
+
+    @override
+    def _load_embeddings(self, index: int) -> List[torch.Tensor]:
+        filename = self.filename(index)
+        embeddings_path = os.path.join(self._root, filename)
+        embeddings = torch.load(embeddings_path, map_location="cpu")
+        if isinstance(embeddings, torch.Tensor):
+            embeddings = [embeddings]
+        return [tensor.squeeze(0) for tensor in embeddings]
+
+    @override
+    def _load_target(self, index: int) -> tv_tensors.Mask:
+        filename = self._data.at[index, self._column_mapping["target"]]
+        mask_path = os.path.join(self._root, filename)
+        semantic_labels = torch.load(mask_path, map_location="cpu")
+        return tv_tensors.Mask(semantic_labels, dtype=torch.int64)  # type: ignore[reportCallIssue]
+
+    @override
+    def __len__(self) -> int:
+        return len(self._data)

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

@@ -0,0 +1,178 @@
+"""LiTS dataset."""
+
+import functools
+import glob
+import os
+from typing import Any, Callable, Dict, List, Literal, Tuple
+
+import torch
+from torchvision import tv_tensors
+from typing_extensions import override
+
+from eva.core import utils
+from eva.vision.data.datasets import _utils as data_utils
+from eva.vision.data.datasets import _validators
+from eva.vision.data.datasets.segmentation import base
+from eva.vision.utils import io
+
+
+class LiTS(base.ImageSegmentation):
+    """LiTS - Liver Tumor Segmentation Challenge.
+
+    Webpage: https://competitions.codalab.org/competitions/17094
+
+    For the splits we follow: https://arxiv.org/pdf/2010.01663v2
+    """
+
+    _train_index_ranges: List[Tuple[int, int]] = [(0, 102)]
+    _val_index_ranges: List[Tuple[int, int]] = [(102, 117)]
+    _test_index_ranges: List[Tuple[int, int]] = [(117, 131)]
+    """Index ranges per split."""
+
+    _sample_every_n_slices: int | None = None
+    """The amount of slices to sub-sample per 3D CT scan image."""
+
+    _expected_dataset_lengths: Dict[str | None, int] = {
+        "train": 39307,
+        "val": 12045,
+        "test": 7286,
+        None: 58638,
+    }
+    """Dataset version and split to the expected size."""
+
+    _license: str = (
+        "Creative Commons Attribution-NonCommercial-NoDerivatives 4.0 International License "
+        "(https://creativecommons.org/licenses/by-nc-nd/4.0/deed.en)"
+    )
+    """Dataset license."""
+
+    def __init__(
+        self,
+        root: str,
+        split: Literal["train", "val", "test"] | None = None,
+        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.
+            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._indices: List[Tuple[int, int]] = []
+
+    @property
+    @override
+    def classes(self) -> List[str]:
+        return ["liver", "tumor"]
+
+    @functools.cached_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_index, _ = self._indices[index]
+        volume_file_path = self._volume_files[sample_index]
+        return os.path.relpath(volume_file_path, self._root)
+
+    @override
+    def configure(self) -> None:
+        self._indices = self._create_indices()
+
+    @override
+    def validate(self) -> None:
+        if len(self._volume_files) != len(self._segmentation_files):
+            raise ValueError(
+                "The number of volume files does not match the number of the segmentation ones."
+            )
+
+        _validators.check_dataset_integrity(
+            self,
+            length=self._expected_dataset_lengths.get(self._split, 0),
+            n_classes=2,
+            first_and_last_labels=("liver", "tumor"),
+        )
+
+    @override
+    def load_image(self, index: int) -> tv_tensors.Image:
+        sample_index, slice_index = self._indices[index]
+        volume_path = self._volume_files[sample_index]
+        image_array = io.read_nifti(volume_path, slice_index)
+        return tv_tensors.Image(image_array.transpose(2, 0, 1))
+
+    @override
+    def load_mask(self, index: int) -> tv_tensors.Mask:
+        sample_index, slice_index = self._indices[index]
+        segmentation_path = self._segmentation_files[sample_index]
+        semantic_labels = io.read_nifti(segmentation_path, slice_index)
+        return tv_tensors.Mask(semantic_labels.squeeze(), dtype=torch.int64)  # type: ignore[reportCallIssue]
+
+    @override
+    def load_metadata(self, index: int) -> Dict[str, Any]:
+        _, slice_index = self._indices[index]
+        return {"slice_index": slice_index}
+
+    @override
+    def __len__(self) -> int:
+        return len(self._indices)
+
+    def _get_number_of_slices_per_volume(self, sample_index: int) -> int:
+        """Returns the total amount of slices of a volume."""
+        file_path = self._volume_files[sample_index]
+        volume_shape = io.fetch_nifti_shape(file_path)
+        return volume_shape[-1]
+
+    @functools.cached_property
+    def _volume_files(self) -> List[str]:
+        files_pattern = os.path.join(self._root, "**", "volume-*.nii")
+        files = glob.glob(files_pattern, recursive=True)
+        return utils.numeric_sort(files)
+
+    @functools.cached_property
+    def _segmentation_files(self) -> List[str]:
+        files_pattern = os.path.join(self._root, "**", "segmentation-*.nii")
+        files = glob.glob(files_pattern, recursive=True)
+        return utils.numeric_sort(files)
+
+    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_volume(sample_idx))
+            if slide_idx % (self._sample_every_n_slices or 1) == 0
+        ]
+        return indices
+
+    def _get_split_indices(self) -> List[int]:
+        """Returns the sample indices for the specified dataset split."""
+        split_index_ranges = {
+            "train": self._train_index_ranges,
+            "val": self._val_index_ranges,
+            "test": self._test_index_ranges,
+            None: [(0, len(self._volume_files))],
+        }
+        index_ranges = split_index_ranges.get(self._split)
+        if index_ranges is None:
+            raise ValueError("Invalid data split. Use 'train', 'val', 'test' or `None`.")
+
+        return data_utils.ranges_to_indices(index_ranges)
+
+    def _print_license(self) -> None:
+        """Prints the dataset license."""
+        print(f"Dataset license: {self._license}")