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

eva/core/data/datasets/classification/__init__.py

@@ -1,5 +1,8 @@
-"""Classification datasets API."""
+"""Embedding cllassification datasets API."""
 
 from eva.core.data.datasets.classification.embeddings import EmbeddingsClassificationDataset
+from eva.core.data.datasets.classification.multi_embeddings import (
+    MultiEmbeddingsClassificationDataset,
+)
 
-__all__ = ["EmbeddingsClassificationDataset"]
+__all__ = ["EmbeddingsClassificationDataset", "MultiEmbeddingsClassificationDataset"]

eva/core/data/datasets/classification/embeddings.py

@@ -1,154 +1,34 @@
 """Embeddings classification dataset."""
 
 import os
-from typing import Callable, Dict, Tuple
 
-import numpy as np
-import pandas as pd
 import torch
 from typing_extensions import override
 
-from eva.core.data.datasets import base
-from eva.core.utils import io
+from eva.core.data.datasets import embeddings as embeddings_base
 
 
-class EmbeddingsClassificationDataset(base.Dataset):
-    """Embeddings classification dataset."""
-
-    default_column_mapping: Dict[str, str] = {
-        "data": "embeddings",
-        "target": "target",
-        "split": "split",
-    }
-    """The default column mapping of the variables to the manifest columns."""
-
-    def __init__(
-        self,
-        root: str,
-        manifest_file: str,
-        split: str | None = None,
-        column_mapping: Dict[str, str] = default_column_mapping,
-        embeddings_transforms: Callable | None = None,
-        target_transforms: Callable | None = None,
-    ) -> None:
-        """Initialize dataset.
-
-        Expects a manifest file listing the paths of .pt files that contain
-        tensor embeddings of shape [embedding_dim] or [1, embedding_dim].
-
-        Args:
-            root: Root directory of the dataset.
-            manifest_file: The path to the manifest file, which is relative to
-                the `root` argument.
-            split: The dataset split to use. The `split` column of the manifest
-                file will be splitted based on this value.
-            column_mapping: Defines the map between the variables and the manifest
-                columns. It will overwrite the `default_column_mapping` with
-                the provided values, so that `column_mapping` can contain only the
-                values which are altered or missing.
-            embeddings_transforms: A function/transform that transforms the embedding.
-            target_transforms: A function/transform that transforms the target.
-        """
-        super().__init__()
-
-        self._root = root
-        self._manifest_file = manifest_file
-        self._split = split
-        self._column_mapping = self.default_column_mapping | column_mapping
-        self._embeddings_transforms = embeddings_transforms
-        self._target_transforms = target_transforms
-
-        self._data: pd.DataFrame
-
-    def filename(self, index: int) -> str:
-        """Returns the filename of the `index`'th data sample.
-
-        Note that this is the relative file path to the root.
-
-        Args:
-            index: The index of the data-sample to select.
-
-        Returns:
-            The filename of the `index`'th data sample.
-        """
-        return self._data.at[index, self._column_mapping["data"]]
+class EmbeddingsClassificationDataset(embeddings_base.EmbeddingsDataset[torch.Tensor]):
+    """Embeddings dataset class for classification tasks."""
 
     @override
-    def setup(self):
-        self._data = self._load_manifest()
-
-    def __getitem__(self, index) -> Tuple[torch.Tensor, np.ndarray]:
-        """Returns the `index`'th data sample.
-
-        Args:
-            index: The index of the data-sample to select.
-
-        Returns:
-            A data sample and its target.
-        """
-        embeddings = self._load_embeddings(index)
-        target = self._load_target(index)
-        return self._apply_transforms(embeddings, target)
-
-    def __len__(self) -> int:
-        """Returns the total length of the data."""
-        return len(self._data)
-
     def _load_embeddings(self, index: int) -> torch.Tensor:
-        """Returns the `index`'th embedding sample.
-
-        Args:
-            index: The index of the data sample to load.
-
-        Returns:
-            The sample embedding as an array.
-        """
         filename = self.filename(index)
         embeddings_path = os.path.join(self._root, filename)
         tensor = torch.load(embeddings_path, map_location="cpu")
+        if isinstance(tensor, list):
+            if len(tensor) > 1:
+                raise ValueError(
+                    f"Expected a single tensor in the .pt file, but found {len(tensor)}."
+                )
+            tensor = tensor[0]
         return tensor.squeeze(0)
 
-    def _load_target(self, index: int) -> np.ndarray:
-        """Returns the `index`'th target sample.
-
-        Args:
-            index: The index of the data sample to load.
-
-        Returns:
-            The sample target as an array.
-        """
+    @override
+    def _load_target(self, index: int) -> torch.Tensor:
         target = self._data.at[index, self._column_mapping["target"]]
-        return np.asarray(target, dtype=np.int64)
-
-    def _load_manifest(self) -> pd.DataFrame:
-        """Loads manifest file and filters the data based on the split column.
-
-        Returns:
-            The data as a pandas DataFrame.
-        """
-        manifest_path = os.path.join(self._root, self._manifest_file)
-        data = io.read_dataframe(manifest_path)
-        if self._split is not None:
-            filtered_data = data.loc[data[self._column_mapping["split"]] == self._split]
-            data = filtered_data.reset_index(drop=True)
-        return data
+        return torch.tensor(target, dtype=torch.int64)
 
-    def _apply_transforms(
-        self, embeddings: torch.Tensor, target: np.ndarray
-    ) -> Tuple[torch.Tensor, np.ndarray]:
-        """Applies the transforms to the provided data and returns them.
-
-        Args:
-            embeddings: The embeddings to be transformed.
-            target: The training target.
-
-        Returns:
-            A tuple with the embeddings and the target transformed.
-        """
-        if self._embeddings_transforms is not None:
-            embeddings = self._embeddings_transforms(embeddings)
-
-        if self._target_transforms is not None:
-            target = self._target_transforms(target)
-
-        return embeddings, target
+    @override
+    def __len__(self) -> int:
+        return len(self._data)

eva/core/data/datasets/classification/multi_embeddings.py

@@ -0,0 +1,110 @@
+"""Dataset class for where a sample corresponds to multiple embeddings."""
+
+import os
+from typing import Callable, Dict, List, Literal
+
+import numpy as np
+import torch
+from typing_extensions import override
+
+from eva.core.data.datasets import embeddings as embeddings_base
+
+
+class MultiEmbeddingsClassificationDataset(embeddings_base.EmbeddingsDataset[torch.Tensor]):
+    """Dataset class for where a sample corresponds to multiple embeddings.
+
+    Example use case: Slide level dataset where each slide has multiple patch embeddings.
+    """
+
+    def __init__(
+        self,
+        root: str,
+        manifest_file: str,
+        split: Literal["train", "val", "test"],
+        column_mapping: Dict[str, str] = embeddings_base.default_column_mapping,
+        embeddings_transforms: Callable | None = None,
+        target_transforms: Callable | None = None,
+    ):
+        """Initialize dataset.
+
+        Expects a manifest file listing the paths of `.pt` files containing tensor embeddings.
+
+        The manifest must have a `column_mapping["multi_id"]` column that contains the
+        unique identifier group of embeddings. For oncology datasets, this would be usually
+        the slide id. Each row in the manifest file points to a .pt file that can contain
+        one or multiple embeddings (either as a list or stacked tensors). There can also be
+        multiple rows for the same `multi_id`, in which case the embeddings from the different
+        .pt files corresponding to that same `multi_id` will be stacked along the first dimension.
+
+        Args:
+            root: Root directory of the dataset.
+            manifest_file: The path to the manifest file, which is relative to
+                the `root` argument.
+            split: The dataset split to use. The `split` column of the manifest
+                file will be splitted based on this value.
+            column_mapping: Defines the map between the variables and the manifest
+                columns. It will overwrite the `default_column_mapping` with
+                the provided values, so that `column_mapping` can contain only the
+                values which are altered or missing.
+            embeddings_transforms: A function/transform that transforms the embedding.
+            target_transforms: A function/transform that transforms the target.
+        """
+        super().__init__(
+            manifest_file=manifest_file,
+            root=root,
+            split=split,
+            column_mapping=column_mapping,
+            embeddings_transforms=embeddings_transforms,
+            target_transforms=target_transforms,
+        )
+
+        self._multi_ids: List[int]
+
+    @override
+    def setup(self):
+        super().setup()
+        self._multi_ids = list(self._data[self._column_mapping["multi_id"]].unique())
+
+    @override
+    def _load_embeddings(self, index: int) -> torch.Tensor:
+        """Loads and stacks all embedding corresponding to the `index`'th multi_id."""
+        # Get all embeddings for the given index (multi_id)
+        multi_id = self._multi_ids[index]
+        embedding_paths = self._data.loc[
+            self._data[self._column_mapping["multi_id"]] == multi_id, self._column_mapping["path"]
+        ].to_list()
+
+        # Load embeddings and stack them accross the first dimension
+        embeddings = []
+        for path in embedding_paths:
+            embedding = torch.load(os.path.join(self._root, path), map_location="cpu")
+            if isinstance(embedding, list):
+                embedding = torch.stack(embedding, dim=0)
+            embeddings.append(embedding.unsqueeze(0) if embedding.ndim == 1 else embedding)
+        embeddings = torch.cat(embeddings, dim=0)
+
+        if not embeddings.ndim == 2:
+            raise ValueError(f"Expected 2D tensor, got {embeddings.ndim} for {multi_id}.")
+
+        return embeddings
+
+    @override
+    def _load_target(self, index: int) -> np.ndarray:
+        """Returns the target corresponding to the `index`'th multi_id.
+
+        This method assumes that all the embeddings corresponding to the same `multi_id`
+        have the same target. If this is not the case, it will raise an error.
+        """
+        multi_id = self._multi_ids[index]
+        targets = self._data.loc[
+            self._data[self._column_mapping["multi_id"]] == multi_id, self._column_mapping["target"]
+        ]
+
+        if not targets.nunique() == 1:
+            raise ValueError(f"Multiple targets found for {multi_id}.")
+
+        return np.asarray(targets.iloc[0], dtype=np.int64)
+
+    @override
+    def __len__(self) -> int:
+        return len(self._multi_ids)

eva/core/data/datasets/embeddings.py

@@ -0,0 +1,167 @@
+"""Base dataset class for Embeddings."""
+
+import abc
+import multiprocessing
+import os
+from typing import Callable, Dict, Generic, Literal, Tuple, TypeVar
+
+import pandas as pd
+import torch
+from typing_extensions import override
+
+from eva.core.data.datasets import base
+from eva.core.utils import io
+
+TargetType = TypeVar("TargetType")
+"""The target data type."""
+
+
+default_column_mapping: Dict[str, str] = {
+    "path": "embeddings",
+    "target": "target",
+    "split": "split",
+    "multi_id": "wsi_id",
+}
+"""The default column mapping of the variables to the manifest columns."""
+
+
+class EmbeddingsDataset(base.Dataset, Generic[TargetType]):
+    """Abstract base class for embedding datasets."""
+
+    def __init__(
+        self,
+        root: str,
+        manifest_file: str,
+        split: Literal["train", "val", "test"] | None = None,
+        column_mapping: Dict[str, str] = default_column_mapping,
+        embeddings_transforms: Callable | None = None,
+        target_transforms: Callable | None = None,
+    ) -> None:
+        """Initialize dataset.
+
+        Expects a manifest file listing the paths of .pt files that contain
+        tensor embeddings of shape [embedding_dim] or [1, embedding_dim].
+
+        Args:
+            root: Root directory of the dataset.
+            manifest_file: The path to the manifest file, which is relative to
+                the `root` argument.
+            split: The dataset split to use. The `split` column of the manifest
+                file will be splitted based on this value.
+            column_mapping: Defines the map between the variables and the manifest
+                columns. It will overwrite the `default_column_mapping` with
+                the provided values, so that `column_mapping` can contain only the
+                values which are altered or missing.
+            embeddings_transforms: A function/transform that transforms the embedding.
+            target_transforms: A function/transform that transforms the target.
+        """
+        super().__init__()
+
+        self._root = root
+        self._manifest_file = manifest_file
+        self._split = split
+        self._column_mapping = default_column_mapping | column_mapping
+        self._embeddings_transforms = embeddings_transforms
+        self._target_transforms = target_transforms
+
+        self._data: pd.DataFrame
+
+        self._set_multiprocessing_start_method()
+
+    def filename(self, index: int) -> str:
+        """Returns the filename of the `index`'th data sample.
+
+        Note that this is the relative file path to the root.
+
+        Args:
+            index: The index of the data-sample to select.
+
+        Returns:
+            The filename of the `index`'th data sample.
+        """
+        return self._data.at[index, self._column_mapping["path"]]
+
+    @override
+    def setup(self):
+        self._data = self._load_manifest()
+
+    @abc.abstractmethod
+    def __len__(self) -> int:
+        """Returns the total length of the data."""
+
+    def __getitem__(self, index) -> Tuple[torch.Tensor, TargetType]:
+        """Returns the `index`'th data sample.
+
+        Args:
+            index: The index of the data-sample to select.
+
+        Returns:
+            A data sample and its target.
+        """
+        embeddings = self._load_embeddings(index)
+        target = self._load_target(index)
+        return self._apply_transforms(embeddings, target)
+
+    @abc.abstractmethod
+    def _load_embeddings(self, index: int) -> torch.Tensor:
+        """Returns the `index`'th embedding sample.
+
+        Args:
+            index: The index of the data sample to load.
+
+        Returns:
+            The embedding sample as a tensor.
+        """
+
+    @abc.abstractmethod
+    def _load_target(self, index: int) -> TargetType:
+        """Returns the `index`'th target sample.
+
+        Args:
+            index: The index of the data sample to load.
+
+        Returns:
+            The sample target as an array.
+        """
+
+    def _load_manifest(self) -> pd.DataFrame:
+        """Loads manifest file and filters the data based on the split column.
+
+        Returns:
+            The data as a pandas DataFrame.
+        """
+        manifest_path = os.path.join(self._root, self._manifest_file)
+        data = io.read_dataframe(manifest_path)
+        if self._split is not None:
+            filtered_data = data.loc[data[self._column_mapping["split"]] == self._split]
+            data = filtered_data.reset_index(drop=True)
+        return data
+
+    def _apply_transforms(
+        self, embeddings: torch.Tensor, target: TargetType
+    ) -> Tuple[torch.Tensor, TargetType]:
+        """Applies the transforms to the provided data and returns them.
+
+        Args:
+            embeddings: The embeddings to be transformed.
+            target: The training target.
+
+        Returns:
+            A tuple with the embeddings and the target transformed.
+        """
+        if self._embeddings_transforms is not None:
+            embeddings = self._embeddings_transforms(embeddings)
+
+        if self._target_transforms is not None:
+            target = self._target_transforms(target)
+
+        return embeddings, target
+
+    def _set_multiprocessing_start_method(self):
+        """Sets the multiprocessing start method to spawn.
+
+        If the start method is not set explicitly, the torch data loaders will
+        use the OS default method, which for some unix systems is `fork` and
+        can lead to runtime issues such as deadlocks in this context.
+        """
+        multiprocessing.set_start_method("spawn", force=True)

eva/core/data/splitting/__init__.py

@@ -0,0 +1,6 @@
+"""Dataset splitting API."""
+
+from eva.core.data.splitting.random import random_split
+from eva.core.data.splitting.stratified import stratified_split
+
+__all__ = ["random_split", "stratified_split"]

eva/core/data/splitting/random.py

@@ -0,0 +1,41 @@
+"""Functions for random splitting."""
+
+from typing import Any, List, Sequence, Tuple
+
+import numpy as np
+
+
+def random_split(
+    samples: Sequence[Any],
+    train_ratio: float,
+    val_ratio: float,
+    test_ratio: float = 0.0,
+    seed: int = 42,
+) -> Tuple[List[int], List[int], List[int] | None]:
+    """Splits the samples into random train, validation, and test (optional) sets.
+
+    Args:
+        samples: The samples to split.
+        train_ratio: The ratio of the training set.
+        val_ratio: The ratio of the validation set.
+        test_ratio: The ratio of the test set (optional).
+        seed: The seed for reproducibility.
+
+    Returns:
+        The indices of the train, validation, and test sets as lists.
+    """
+    if train_ratio + val_ratio + (test_ratio or 0) != 1:
+        raise ValueError("The sum of the ratios must be equal to 1.")
+
+    np.random.seed(seed)
+    n_samples = len(samples)
+    indices = np.random.permutation(n_samples)
+
+    n_train = int(np.floor(train_ratio * n_samples))
+    n_val = n_samples - n_train if test_ratio == 0.0 else int(np.floor(val_ratio * n_samples)) or 1
+
+    train_indices = list(indices[:n_train])
+    val_indices = list(indices[n_train : n_train + n_val])
+    test_indices = list(indices[n_train + n_val :]) if test_ratio > 0.0 else None
+
+    return train_indices, val_indices, test_indices

eva/core/data/splitting/stratified.py

@@ -0,0 +1,56 @@
+"""Functions for stratified splitting."""
+
+from typing import Any, List, Sequence, Tuple
+
+import numpy as np
+
+
+def stratified_split(
+    samples: Sequence[Any],
+    targets: Sequence[Any],
+    train_ratio: float,
+    val_ratio: float,
+    test_ratio: float = 0.0,
+    seed: int = 42,
+) -> Tuple[List[int], List[int], List[int] | None]:
+    """Splits the samples into stratified train, validation, and test (optional) sets.
+
+    Args:
+        samples: The samples to split.
+        targets: The corresponding targets used for stratification.
+        train_ratio: The ratio of the training set.
+        val_ratio: The ratio of the validation set.
+        test_ratio: The ratio of the test set (optional).
+        seed: The seed for reproducibility.
+
+    Returns:
+        The indices of the train, validation, and test sets.
+    """
+    if len(samples) != len(targets):
+        raise ValueError("The number of samples and targets must be equal.")
+    if train_ratio + val_ratio + (test_ratio or 0) != 1:
+        raise ValueError("The sum of the ratios must be equal to 1.")
+
+    np.random.seed(seed)
+    unique_classes, y_indices = np.unique(targets, return_inverse=True)
+    n_classes = unique_classes.shape[0]
+
+    train_indices, val_indices, test_indices = [], [], []
+
+    for c in range(n_classes):
+        class_indices = np.where(y_indices == c)[0]
+        np.random.shuffle(class_indices)
+
+        n_train = int(np.floor(train_ratio * len(class_indices))) or 1
+        n_val = (
+            len(class_indices) - n_train
+            if test_ratio == 0.0
+            else int(np.floor(val_ratio * len(class_indices))) or 1
+        )
+
+        train_indices.extend(class_indices[:n_train])
+        val_indices.extend(class_indices[n_train : n_train + n_val])
+        if test_ratio > 0.0:
+            test_indices.extend(class_indices[n_train + n_val :])
+
+    return train_indices, val_indices, test_indices or None

eva/core/data/transforms/__init__.py

@@ -1,5 +1,7 @@
 """Core data transforms."""
 
 from eva.core.data.transforms.dtype import ArrayToFloatTensor, ArrayToTensor
+from eva.core.data.transforms.padding import Pad2DTensor
+from eva.core.data.transforms.sampling import SampleFromAxis
 
-__all__ = ["ArrayToFloatTensor", "ArrayToTensor"]
+__all__ = ["ArrayToFloatTensor", "ArrayToTensor", "Pad2DTensor", "SampleFromAxis"]

eva/core/data/transforms/padding/__init__.py

@@ -0,0 +1,5 @@
+"""Padding related transformations."""
+
+from eva.core.data.transforms.padding.pad_2d_tensor import Pad2DTensor
+
+__all__ = ["Pad2DTensor"]

eva/core/data/transforms/padding/pad_2d_tensor.py

@@ -0,0 +1,38 @@
+"""Padding transformation for 2D tensors."""
+
+import torch
+import torch.nn.functional
+
+
+class Pad2DTensor:
+    """Pads a 2D tensor to a fixed dimension accross the first dimension."""
+
+    def __init__(self, pad_size: int, pad_value: int | float = float("-inf")):
+        """Initialize the transformation.
+
+        Args:
+            pad_size: The size to pad the tensor to. If the tensor is larger than this size,
+                no padding will be applied.
+            pad_value: The value to use for padding.
+        """
+        self._pad_size = pad_size
+        self._pad_value = pad_value
+
+    def __call__(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Call method for the transformation.
+
+        Args:
+            tensor: The input tensor of shape [n, embedding_dim].
+
+        Returns:
+            A tensor of shape [max(n, pad_dim), embedding_dim].
+        """
+        n_pad_values = self._pad_size - tensor.size(0)
+        if n_pad_values > 0:
+            tensor = torch.nn.functional.pad(
+                tensor,
+                pad=(0, 0, 0, n_pad_values),
+                mode="constant",
+                value=self._pad_value,
+            )
+        return tensor

eva/core/data/transforms/sampling/__init__.py

@@ -0,0 +1,5 @@
+"""Sampling related transformations."""
+
+from eva.core.data.transforms.sampling.sample_from_axis import SampleFromAxis
+
+__all__ = ["SampleFromAxis"]

eva/core/data/transforms/sampling/sample_from_axis.py

@@ -0,0 +1,40 @@
+"""Sampling transformations."""
+
+import torch
+
+
+class SampleFromAxis:
+    """Samples n_samples entries from a tensor along a given axis."""
+
+    def __init__(self, n_samples: int, seed: int = 42, axis: int = 0):
+        """Initialize the transformation.
+
+        Args:
+            n_samples: The number of samples to draw.
+            seed: The seed to use for sampling.
+            axis: The axis along which to sample.
+        """
+        self._seed = seed
+        self._n_samples = n_samples
+        self._axis = axis
+        self._generator = self._get_generator()
+
+    def _get_generator(self):
+        """Return a torch random generator with fixed seed."""
+        generator = torch.Generator()
+        generator.manual_seed(self._seed)
+        return generator
+
+    def __call__(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Call method for the transformation.
+
+        Args:
+            tensor: The input tensor of shape [n, embedding_dim].
+
+        Returns:
+            A tensor of shape [n_samples, embedding_dim].
+        """
+        indices = torch.randperm(tensor.size(self._axis), generator=self._generator)[
+            : self._n_samples
+        ]
+        return tensor.index_select(self._axis, indices)

eva/core/loggers/__init__.py

@@ -0,0 +1,7 @@
+"""Experimental loggers API."""
+
+from eva.core.loggers.dummy import DummyLogger
+from eva.core.loggers.experimental_loggers import ExperimentalLoggers
+from eva.core.loggers.log import log_parameters
+
+__all__ = ["DummyLogger", "ExperimentalLoggers", "log_parameters"]