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

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

@@ -0,0 +1,105 @@
+"""WSI classification dataset."""
+
+import os
+from typing import Any, Callable, Dict, Literal, Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+from torchvision import tv_tensors
+from typing_extensions import override
+
+from eva.vision.data.datasets import wsi
+from eva.vision.data.datasets.classification import base
+from eva.vision.data.wsi.patching import samplers
+
+
+class WsiClassificationDataset(wsi.MultiWsiDataset, base.ImageClassification):
+    """A general dataset class for whole-slide image classification using manifest files."""
+
+    default_column_mapping: Dict[str, str] = {
+        "path": "path",
+        "target": "target",
+        "split": "split",
+    }
+
+    def __init__(
+        self,
+        root: str,
+        manifest_file: str,
+        width: int,
+        height: int,
+        target_mpp: float,
+        sampler: samplers.Sampler,
+        backend: str = "openslide",
+        split: Literal["train", "val", "test"] | None = None,
+        image_transforms: Callable | None = None,
+        column_mapping: Dict[str, str] = default_column_mapping,
+    ):
+        """Initializes the dataset.
+
+        Args:
+            root: Root directory of the dataset.
+            manifest_file: The path to the manifest file, relative to
+                the `root` argument. The `path` column is expected to contain
+                relative paths to the whole-slide images.
+            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.
+            sampler: The sampler to use for sampling patch coordinates.
+            backend: The backend to use for reading the whole-slide images.
+            split: The split of the dataset to load.
+            image_transforms: Transforms to apply to the extracted image patches.
+            column_mapping: Mapping of the columns in the manifest file.
+        """
+        self._split = split
+        self._column_mapping = self.default_column_mapping | column_mapping
+        self._manifest = self._load_manifest(os.path.join(root, manifest_file))
+
+        wsi.MultiWsiDataset.__init__(
+            self,
+            root=root,
+            file_paths=self._manifest[self._column_mapping["path"]].tolist(),
+            width=width,
+            height=height,
+            sampler=sampler,
+            target_mpp=target_mpp,
+            backend=backend,
+            image_transforms=image_transforms,
+        )
+
+    @override
+    def filename(self, index: int) -> str:
+        path = self._manifest.at[self._get_dataset_idx(index), self._column_mapping["path"]]
+        return os.path.basename(path) if os.path.isabs(path) else path
+
+    @override
+    def __getitem__(self, index: int) -> Tuple[tv_tensors.Image, torch.Tensor, Dict[str, Any]]:
+        return base.ImageClassification.__getitem__(self, index)
+
+    @override
+    def load_image(self, index: int) -> tv_tensors.Image:
+        return wsi.MultiWsiDataset.__getitem__(self, index)
+
+    @override
+    def load_target(self, index: int) -> np.ndarray:
+        target = self._manifest.at[self._get_dataset_idx(index), self._column_mapping["target"]]
+        return np.asarray(target)
+
+    @override
+    def load_metadata(self, index: int) -> Dict[str, Any]:
+        return {"wsi_id": self.filename(index).split(".")[0]}
+
+    def _load_manifest(self, manifest_path: str) -> pd.DataFrame:
+        df = pd.read_csv(manifest_path)
+
+        missing_columns = set(self._column_mapping.values()) - set(df.columns)
+        if self._split is None:
+            missing_columns = missing_columns - {self._column_mapping["split"]}
+        if missing_columns:
+            raise ValueError(f"Missing columns in the manifest file: {missing_columns}")
+
+        if self._split is not None:
+            df = df.loc[df[self._column_mapping["split"]] == self._split]
+
+        return df.reset_index(drop=True)

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

@@ -1,6 +1,19 @@
 """Segmentation datasets API."""
 
 from eva.vision.data.datasets.segmentation.base import ImageSegmentation
-from eva.vision.data.datasets.segmentation.total_segmentator import TotalSegmentator2D
+from eva.vision.data.datasets.segmentation.bcss import BCSS
+from eva.vision.data.datasets.segmentation.consep import CoNSeP
+from eva.vision.data.datasets.segmentation.embeddings import EmbeddingsSegmentationDataset
+from eva.vision.data.datasets.segmentation.lits import LiTS
+from eva.vision.data.datasets.segmentation.monusac import MoNuSAC
+from eva.vision.data.datasets.segmentation.total_segmentator_2d import TotalSegmentator2D
 
-__all__ = ["ImageSegmentation", "TotalSegmentator2D"]
+__all__ = [
+    "ImageSegmentation",
+    "BCSS",
+    "CoNSeP",
+    "EmbeddingsSegmentationDataset",
+    "LiTS",
+    "MoNuSAC",
+    "TotalSegmentator2D",
+]

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

@@ -0,0 +1,38 @@
+from typing import Any, Tuple
+
+import numpy.typing as npt
+
+from eva.vision.data.datasets import wsi
+
+
+def get_coords_at_index(
+    dataset: wsi.MultiWsiDataset, index: int
+) -> Tuple[Tuple[int, int], int, int]:
+    """Returns the coordinates ((x,y),width,height) of the patch at the given index.
+
+    Args:
+        dataset: The WSI dataset instance.
+        index: The sample index.
+    """
+    image_index = dataset._get_dataset_idx(index)
+    patch_index = index if image_index == 0 else index - dataset.cumulative_sizes[image_index - 1]
+    wsi_dataset = dataset.datasets[image_index]
+    if isinstance(wsi_dataset, wsi.WsiDataset):
+        coords = wsi_dataset._coords
+        return coords.x_y[patch_index], coords.width, coords.height
+    else:
+        raise Exception(f"Expected WsiDataset, got {type(wsi_dataset)}")
+
+
+def extract_mask_patch(
+    mask: npt.NDArray[Any], dataset: wsi.MultiWsiDataset, index: int
+) -> npt.NDArray[Any]:
+    """Reads the mask patch at the coordinates corresponding to the dataset index.
+
+    Args:
+        mask: The mask array.
+        dataset: The WSI dataset instance.
+        index: The sample index.
+    """
+    (x, y), width, height = get_coords_at_index(dataset, index)
+    return mask[y : y + height, x : x + width]

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

@@ -12,10 +12,7 @@ from eva.vision.data.datasets import vision
 class ImageSegmentation(vision.VisionDataset[Tuple[tv_tensors.Image, tv_tensors.Mask]], abc.ABC):
     """Image segmentation abstract dataset."""
 
-    def __init__(
-        self,
-        transforms: Callable | None = None,
-    ) -> None:
+    def __init__(self, transforms: Callable | None = None) -> None:
         """Initializes the image segmentation base class.
 
         Args:
@@ -34,17 +31,6 @@ class ImageSegmentation(vision.VisionDataset[Tuple[tv_tensors.Image, tv_tensors.
     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) -> tv_tensors.Image:
         """Loads and returns the `index`'th image sample.
@@ -68,16 +54,29 @@ class ImageSegmentation(vision.VisionDataset[Tuple[tv_tensors.Image, tv_tensors.
             values which represent the pixel class id.
         """
 
+    def load_metadata(self, index: int) -> 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
     @override
     def __len__(self) -> int:
         raise NotImplementedError
 
     @override
-    def __getitem__(self, index: int) -> Tuple[tv_tensors.Image, tv_tensors.Mask]:
+    def __getitem__(self, index: int) -> Tuple[tv_tensors.Image, tv_tensors.Mask, Dict[str, Any]]:
         image = self.load_image(index)
         mask = self.load_mask(index)
-        return self._apply_transforms(image, mask)
+        metadata = self.load_metadata(index) or {}
+        image_tensor, mask_tensor = self._apply_transforms(image, mask)
+        return image_tensor, mask_tensor, metadata
 
     def _apply_transforms(
         self, image: tv_tensors.Image, mask: tv_tensors.Mask

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)