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

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

@@ -0,0 +1,184 @@
+"""PANDA dataset class."""
+
+import functools
+import glob
+import os
+from typing import Any, Callable, Dict, List, Literal, Tuple
+
+import pandas as pd
+import torch
+from torchvision import tv_tensors
+from torchvision.datasets import utils
+from torchvision.transforms.v2 import functional
+from typing_extensions import override
+
+from eva.core.data import splitting
+from eva.vision.data.datasets import _validators, structs, wsi
+from eva.vision.data.datasets.classification import base
+from eva.vision.data.wsi.patching import samplers
+
+
+class PANDA(wsi.MultiWsiDataset, base.ImageClassification):
+    """Dataset class for PANDA images and corresponding targets."""
+
+    _train_split_ratio: float = 0.7
+    """Train split ratio."""
+
+    _val_split_ratio: float = 0.15
+    """Validation split ratio."""
+
+    _test_split_ratio: float = 0.15
+    """Test split ratio."""
+
+    _resources: List[structs.DownloadResource] = [
+        structs.DownloadResource(
+            filename="train_with_noisy_labels.csv",
+            url="https://raw.githubusercontent.com/analokmaus/kaggle-panda-challenge-public/master/train.csv",
+            md5="5e4bfc78bda9603d2e2faf3ed4b21dfa",
+        )
+    ]
+    """Download resources."""
+
+    def __init__(
+        self,
+        root: str,
+        sampler: samplers.Sampler,
+        split: Literal["train", "val", "test"] | None = None,
+        width: int = 224,
+        height: int = 224,
+        target_mpp: float = 0.5,
+        backend: str = "openslide",
+        image_transforms: Callable | None = None,
+        seed: int = 42,
+    ) -> None:
+        """Initializes the dataset.
+
+        Args:
+            root: Root directory of the dataset.
+            sampler: The sampler to use for sampling patch coordinates.
+            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.
+            image_transforms: Transforms to apply to the extracted image patches.
+            seed: Random seed for reproducibility.
+        """
+        self._split = split
+        self._root = root
+        self._seed = seed
+
+        self._download_resources()
+
+        wsi.MultiWsiDataset.__init__(
+            self,
+            root=root,
+            file_paths=self._load_file_paths(split),
+            width=width,
+            height=height,
+            sampler=sampler,
+            target_mpp=target_mpp,
+            backend=backend,
+            image_transforms=image_transforms,
+        )
+
+    @property
+    @override
+    def classes(self) -> List[str]:
+        return ["0", "1", "2", "3", "4", "5"]
+
+    @functools.cached_property
+    def annotations(self) -> pd.DataFrame:
+        """Loads the dataset labels."""
+        path = os.path.join(self._root, "train_with_noisy_labels.csv")
+        return pd.read_csv(path, index_col="image_id")
+
+    @override
+    def prepare_data(self) -> None:
+        _validators.check_dataset_exists(self._root, False)
+
+        if not os.path.isdir(os.path.join(self._root, "train_images")):
+            raise FileNotFoundError("'train_images' directory not found in the root folder.")
+        if not os.path.isfile(os.path.join(self._root, "train_with_noisy_labels.csv")):
+            raise FileNotFoundError("'train.csv' file not found in the root folder.")
+
+    def _download_resources(self) -> None:
+        """Downloads the dataset resources."""
+        for resource in self._resources:
+            utils.download_url(resource.url, self._root, resource.filename, resource.md5)
+
+    @override
+    def validate(self) -> None:
+        _validators.check_dataset_integrity(
+            self,
+            length=None,
+            n_classes=6,
+            first_and_last_labels=("0", "5"),
+        )
+
+    @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:
+        image_array = wsi.MultiWsiDataset.__getitem__(self, index)
+        return functional.to_image(image_array)
+
+    @override
+    def load_target(self, index: int) -> torch.Tensor:
+        file_path = self._file_paths[self._get_dataset_idx(index)]
+        return torch.tensor(self._get_target_from_path(file_path), dtype=torch.int64)
+
+    @override
+    def load_metadata(self, index: int) -> Dict[str, Any]:
+        return {"wsi_id": self.filename(index).split(".")[0]}
+
+    def _load_file_paths(self, split: Literal["train", "val", "test"] | None = None) -> List[str]:
+        """Loads the file paths of the corresponding dataset split."""
+        image_dir = os.path.join(self._root, "train_images")
+        file_paths = sorted(glob.glob(os.path.join(image_dir, "*.tiff")))
+        file_paths = [os.path.relpath(path, self._root) for path in file_paths]
+        if len(file_paths) != len(self.annotations):
+            raise ValueError(
+                f"Expected {len(self.annotations)} images, found {len(file_paths)} in {image_dir}."
+            )
+        file_paths = self._filter_noisy_labels(file_paths)
+        targets = [self._get_target_from_path(file_path) for file_path in file_paths]
+
+        train_indices, val_indices, test_indices = splitting.stratified_split(
+            samples=file_paths,
+            targets=targets,
+            train_ratio=self._train_split_ratio,
+            val_ratio=self._val_split_ratio,
+            test_ratio=self._test_split_ratio,
+            seed=self._seed,
+        )
+
+        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 "test":
+                return [file_paths[i] for i in test_indices or []]
+            case None:
+                return file_paths
+            case _:
+                raise ValueError("Invalid split. Use 'train', 'val', 'test' or `None`.")
+
+    def _filter_noisy_labels(self, file_paths: List[str]):
+        is_noisy_filter = self.annotations["noise_ratio_10"] == 0
+        non_noisy_image_ids = set(self.annotations.loc[~is_noisy_filter].index)
+        filtered_file_paths = [
+            file_path
+            for file_path in file_paths
+            if self._get_id_from_path(file_path) in non_noisy_image_ids
+        ]
+        return filtered_file_paths
+
+    def _get_target_from_path(self, file_path: str) -> int:
+        return self.annotations.loc[self._get_id_from_path(file_path), "isup_grade"]
+
+    def _get_id_from_path(self, file_path: str) -> str:
+        return os.path.basename(file_path).replace(".tiff", "")

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

@@ -4,8 +4,10 @@ import os
 from typing import Callable, Dict, List, Literal
 
 import h5py
-import numpy as np
+import torch
+from torchvision import tv_tensors
 from torchvision.datasets import utils
+from torchvision.transforms.v2 import functional
 from typing_extensions import override
 
 from eva.vision.data.datasets import _validators, structs
@@ -70,8 +72,7 @@ class PatchCamelyon(base.ImageClassification):
         root: str,
         split: Literal["train", "val", "test"],
         download: bool = False,
-        image_transforms: Callable | None = None,
-        target_transforms: Callable | None = None,
+        transforms: Callable | None = None,
     ) -> None:
         """Initializes the dataset.
 
@@ -82,15 +83,10 @@ class PatchCamelyon(base.ImageClassification):
             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.
+            transforms: A function/transform which returns a transformed
+                version of the raw data samples.
         """
-        super().__init__(
-            image_transforms=image_transforms,
-            target_transforms=target_transforms,
-        )
+        super().__init__(transforms=transforms)
 
         self._root = root
         self._split = split
@@ -131,13 +127,13 @@ class PatchCamelyon(base.ImageClassification):
         )
 
     @override
-    def load_image(self, index: int) -> np.ndarray:
+    def load_image(self, index: int) -> tv_tensors.Image:
         return self._load_from_h5("x", index)
 
     @override
-    def load_target(self, index: int) -> np.ndarray:
+    def load_target(self, index: int) -> torch.Tensor:
         target = self._load_from_h5("y", index).squeeze()
-        return np.asarray(target, dtype=np.int64)
+        return torch.tensor(target, dtype=torch.float32)
 
     @override
     def __len__(self) -> int:
@@ -162,7 +158,7 @@ class PatchCamelyon(base.ImageClassification):
         self,
         data_key: Literal["x", "y"],
         index: int | None = None,
-    ) -> np.ndarray:
+    ) -> tv_tensors.Image:
         """Load data or targets from an HDF5 file.
 
         Args:
@@ -176,7 +172,8 @@ class PatchCamelyon(base.ImageClassification):
         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
+            image_array = data[:] if index is None else data[index]  # type: ignore
+            return functional.to_image(image_array)  # type: ignore
 
     def _fetch_dataset_length(self) -> int:
         """Fetches the dataset split length from its HDF5 file."""

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

@@ -3,38 +3,25 @@
 import abc
 from typing import Any, Callable, Dict, List, Tuple
 
-import numpy as np
+from torchvision import tv_tensors
 from typing_extensions import override
 
 from eva.vision.data.datasets import vision
 
 
-class ImageSegmentation(vision.VisionDataset[Tuple[np.ndarray, np.ndarray]], abc.ABC):
+class ImageSegmentation(vision.VisionDataset[Tuple[tv_tensors.Image, tv_tensors.Mask]], 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:
+    def __init__(self, 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
+            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
+        self._transforms = transforms
 
     @property
     def classes(self) -> List[str] | None:
@@ -44,37 +31,38 @@ class ImageSegmentation(vision.VisionDataset[Tuple[np.ndarray, np.ndarray]], abc
     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.
+    @abc.abstractmethod
+    def load_image(self, index: int) -> tv_tensors.Image:
+        """Loads and returns the `index`'th image sample.
 
         Args:
-            index: The index of the data sample to return the metadata of.
-                If `None`, it will return the metadata of the current dataset.
+            index: The index of the data sample to load.
 
         Returns:
-            The sample metadata.
+            An image torchvision tensor (channels, height, width).
         """
 
     @abc.abstractmethod
-    def load_image(self, index: int) -> np.ndarray:
-        """Loads and returns the `index`'th image sample.
+    def load_mask(self, index: int) -> tv_tensors.Mask:
+        """Returns the `index`'th target masks sample.
 
         Args:
-            index: The index of the data sample to load.
+            index: The index of the data sample target masks to load.
 
         Returns:
-            The image as a numpy array.
+            The semantic mask as a (H x W) shaped tensor with integer
+            values which represent the pixel class id.
         """
 
-    @abc.abstractmethod
-    def load_mask(self, index: int) -> np.ndarray:
-        """Returns the `index`'th target mask sample.
+    def load_metadata(self, index: int) -> Dict[str, Any] | None:
+        """Returns the dataset metadata.
 
         Args:
-            index: The index of the data sample target mask to load.
+            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 mask as a stack of binary mask arrays (label, height, width).
+            The sample metadata.
         """
 
     @abc.abstractmethod
@@ -83,30 +71,26 @@ class ImageSegmentation(vision.VisionDataset[Tuple[np.ndarray, np.ndarray]], abc
         raise NotImplementedError
 
     @override
-    def __getitem__(self, index: int) -> Tuple[np.ndarray, np.ndarray]:
+    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: np.ndarray, target: np.ndarray
-    ) -> Tuple[np.ndarray, np.ndarray]:
+        self, image: tv_tensors.Image, mask: tv_tensors.Mask
+    ) -> Tuple[tv_tensors.Image, tv_tensors.Mask]:
         """Applies the transforms to the provided data and returns them.
 
         Args:
             image: The desired image.
-            target: The target of the image.
+            mask: The target segmentation mask.
 
         Returns:
-            A tuple with the image and the target transformed.
+            A tuple with the image and the masks 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)
+        if self._transforms is not None:
+            image, mask = self._transforms(image, mask)
 
-        return image, target
+        return image, mask