kaiko-eva 0.0.0.dev8__py3-none-any.whl → 0.0.2__py3-none-any.whl

eva/core/callbacks/__init__.py

@@ -1,5 +1,6 @@
 """Callbacks API."""
 
+from eva.core.callbacks.config import ConfigurationLogger
 from eva.core.callbacks.writers import EmbeddingsWriter
 
-__all__ = ["EmbeddingsWriter"]
+__all__ = ["ConfigurationLogger", "EmbeddingsWriter"]

eva/core/callbacks/config.py

@@ -0,0 +1,143 @@
+"""Configuration logger callback."""
+
+import ast
+import os
+import sys
+from types import BuiltinFunctionType
+from typing import Any, Dict, List
+
+import lightning.pytorch as pl
+import yaml
+from lightning_fabric.utilities import cloud_io
+from loguru import logger as cli_logger
+from omegaconf import OmegaConf
+from typing_extensions import TypeGuard, override
+
+from eva.core import loggers
+
+
+class ConfigurationLogger(pl.Callback):
+    """Logs the submitted configuration to the experimental logger."""
+
+    _save_as: str = "config.yaml"
+
+    def __init__(self, verbose: bool = True) -> None:
+        """Initializes the callback.
+
+        Args:
+            verbose: Whether to print the configurations to print the
+                configuration to the terminal.
+        """
+        super().__init__()
+
+        self._verbose = verbose
+
+    @override
+    def setup(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        stage: str | None = None,
+    ) -> None:
+        log_dir = trainer.log_dir
+        if not _logdir_exists(log_dir):
+            return
+
+        configuration = _load_submitted_config()
+
+        if self._verbose:
+            config_as_text = yaml.dump(configuration, sort_keys=False)
+            print(f"Configuration:\033[94m\n---\n{config_as_text}\033[0m")
+
+        save_as = os.path.join(log_dir, self._save_as)
+        fs = cloud_io.get_filesystem(log_dir)
+        with fs.open(save_as, "w") as output_file:
+            yaml.dump(configuration, output_file, sort_keys=False)
+
+        loggers.log_parameters(trainer.loggers, tag="configuration", parameters=configuration)
+
+
+def _logdir_exists(logdir: str | None, verbose: bool = True) -> TypeGuard[str]:
+    """Checks if the trainer has a log directory.
+
+    Args:
+        logdir: Trainer's logdir.
+        name: The name to log with.
+        verbose: Whether to log if it does not exist.
+
+    Returns:
+        A bool indicating if the log directory exists or not.
+    """
+    exists = isinstance(logdir, str)
+    if not exists and verbose:
+        print("\n")
+        cli_logger.warning("Log directory is `None`. Configuration file will not be logged.\n")
+    return exists
+
+
+def _load_submitted_config() -> Dict[str, Any]:
+    """Retrieves and loads the submitted configuration.
+
+    Returns:
+        The path to the configuration file.
+    """
+    config_paths = _fetch_submitted_config_path()
+    return _load_yaml_files(config_paths)
+
+
+def _fetch_submitted_config_path() -> List[str]:
+    """Fetches the config path from command line arguments.
+
+    Returns:
+        The path to the configuration file.
+    """
+    return list(filter(lambda f: f.endswith(".yaml"), sys.argv))
+
+
+def _load_yaml_files(paths: List[str]) -> Dict[str, Any]:
+    """Loads yaml files and merge them from multiple paths.
+
+    Args:
+        paths: The paths to the yaml files.
+
+    Returns:
+        The merged configurations as a dictionary.
+    """
+    merged_config = {}
+    for config_path in paths:
+        fs = cloud_io.get_filesystem(config_path)
+        with fs.open(config_path, "r") as file:
+            omegaconf_file = OmegaConf.load(file)  # type: ignore
+            config_dict = OmegaConf.to_object(omegaconf_file)  # type: ignore
+            parsed_config = _type_resolver(config_dict)  # type: ignore
+            merged_config.update(parsed_config)
+    return merged_config
+
+
+def _type_resolver(mapping: Dict[str, Any]) -> Dict[str, Any]:
+    """Parses the string values of a dictionary in-place.
+
+    Args:
+        mapping: A dictionary object.
+
+    Returns:
+        The mapping with the formatted values.
+    """
+    for key, value in mapping.items():
+        if isinstance(value, dict):
+            formatted_value = _type_resolver(value)
+        elif isinstance(value, list) and isinstance(value[0], dict):
+            formatted_value = [_type_resolver(subvalue) for subvalue in value]
+        else:
+            try:
+                parsed_value = ast.literal_eval(value)  # type: ignore
+                formatted_value = (
+                    value if isinstance(parsed_value, BuiltinFunctionType) else parsed_value
+                )
+
+            except Exception:
+                formatted_value = value
+
+        mapping[key] = formatted_value
+
+    return mapping

eva/core/data/datasets/__init__.py

@@ -1,7 +1,15 @@
 """Datasets API."""
 
 from eva.core.data.datasets.base import Dataset
-from eva.core.data.datasets.classification import EmbeddingsClassificationDataset
 from eva.core.data.datasets.dataset import TorchDataset
+from eva.core.data.datasets.embeddings import (
+    EmbeddingsClassificationDataset,
+    MultiEmbeddingsClassificationDataset,
+)
 
-__all__ = ["Dataset", "EmbeddingsClassificationDataset", "TorchDataset"]
+__all__ = [
+    "Dataset",
+    "EmbeddingsClassificationDataset",
+    "MultiEmbeddingsClassificationDataset",
+    "TorchDataset",
+]

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

@@ -0,0 +1,13 @@
+"""Datasets API."""
+
+from eva.core.data.datasets.embeddings.base import EmbeddingsDataset
+from eva.core.data.datasets.embeddings.classification import (
+    EmbeddingsClassificationDataset,
+    MultiEmbeddingsClassificationDataset,
+)
+
+__all__ = [
+    "EmbeddingsDataset",
+    "EmbeddingsClassificationDataset",
+    "MultiEmbeddingsClassificationDataset",
+]

eva/core/data/datasets/{classification/embeddings.py → embeddings/base.py}

@@ -1,7 +1,8 @@
-"""Embeddings classification dataset."""
+"""Base dataset class for Embeddings."""
 
+import abc
 import os
-from typing import Callable, Dict, Tuple
+from typing import Callable, Dict, Literal, Tuple
 
 import numpy as np
 import pandas as pd
@@ -11,22 +12,23 @@ from typing_extensions import override
 from eva.core.data.datasets import base
 from eva.core.utils import io
 
+default_column_mapping: Dict[str, str] = {
+    "path": "embeddings",
+    "target": "target",
+    "split": "split",
+    "multi_id": "slide_id",
+}
+"""The default column mapping of the variables to the manifest columns."""
 
-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."""
+class EmbeddingsDataset(base.Dataset):
+    """Abstract base class for embedding datasets."""
 
     def __init__(
         self,
         root: str,
         manifest_file: str,
-        split: str | None = None,
+        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,
@@ -54,12 +56,38 @@ class EmbeddingsClassificationDataset(base.Dataset):
         self._root = root
         self._manifest_file = manifest_file
         self._split = split
-        self._column_mapping = self.default_column_mapping | column_mapping
+        self._column_mapping = default_column_mapping | column_mapping
         self._embeddings_transforms = embeddings_transforms
         self._target_transforms = target_transforms
 
         self._data: pd.DataFrame
 
+    @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) -> 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.
+        """
+
+    @abc.abstractmethod
+    def __len__(self) -> int:
+        """Returns the total length of the data."""
+
     def filename(self, index: int) -> str:
         """Returns the filename of the `index`'th data sample.
 
@@ -71,7 +99,7 @@ class EmbeddingsClassificationDataset(base.Dataset):
         Returns:
             The filename of the `index`'th data sample.
         """
-        return self._data.at[index, self._column_mapping["data"]]
+        return self._data.at[index, self._column_mapping["path"]]
 
     @override
     def setup(self):
@@ -90,36 +118,6 @@ class EmbeddingsClassificationDataset(base.Dataset):
         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")
-        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.
-        """
-        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.
 

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

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

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

@@ -0,0 +1,66 @@
+"""Embeddings classification dataset."""
+
+import os
+from typing import Callable, Dict, Literal
+
+import numpy as np
+import torch
+from typing_extensions import override
+
+from eva.core.data.datasets.embeddings import base
+
+
+class EmbeddingsClassificationDataset(base.EmbeddingsDataset):
+    """Embeddings dataset class for classification tasks."""
+
+    def __init__(
+        self,
+        root: str,
+        manifest_file: str,
+        split: Literal["train", "val", "test"] | None = None,
+        column_mapping: Dict[str, str] = base.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__(
+            root=root,
+            manifest_file=manifest_file,
+            split=split,
+            column_mapping=column_mapping,
+            embeddings_transforms=embeddings_transforms,
+            target_transforms=target_transforms,
+        )
+
+    @override
+    def _load_embeddings(self, index: int) -> torch.Tensor:
+        filename = self.filename(index)
+        embeddings_path = os.path.join(self._root, filename)
+        tensor = torch.load(embeddings_path, map_location="cpu")
+        return tensor.squeeze(0)
+
+    @override
+    def _load_target(self, index: int) -> np.ndarray:
+        target = self._data.at[index, self._column_mapping["target"]]
+        return np.asarray(target, dtype=np.int64)
+
+    @override
+    def __len__(self) -> int:
+        return len(self._data)

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

@@ -0,0 +1,106 @@
+"""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.embeddings import base
+
+
+class MultiEmbeddingsClassificationDataset(base.EmbeddingsDataset):
+    """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] = 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. 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()
+        embedding_paths = [os.path.join(self._root, path) for path in embedding_paths]
+
+        # Load embeddings and stack them accross the first dimension
+        embeddings = [torch.load(path, map_location="cpu") for path in embedding_paths]
+        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._data)

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"]

eva/core/loggers/dummy.py

@@ -0,0 +1,38 @@
+"""Dummy logger class."""
+
+import lightning.pytorch.loggers.logger
+
+
+class DummyLogger(lightning.pytorch.loggers.logger.DummyLogger):
+    """Dummy logger class.
+
+    This logger is currently used as a placeholder when saving results
+    to remote storage, as common lightning loggers do not work
+    with azure blob storage:
+
+    <https://github.com/Lightning-AI/pytorch-lightning/issues/18861>
+    <https://github.com/Lightning-AI/pytorch-lightning/issues/19736>
+
+    Simply disabling the loggers when pointing to remote storage doesn't work
+    because callbacks such as LearningRateMonitor or ModelCheckpoint require a
+    logger to be present.
+    """
+
+    def __init__(self, save_dir: str) -> None:
+        """Initializes the logger.
+
+        Args:
+            save_dir: The save directory (this logger does not save anything,
+                but callbacks might use this path to save their outputs).
+        """
+        super().__init__()
+        self._save_dir = save_dir
+
+    @property
+    def save_dir(self) -> str:
+        """Returns the save directory."""
+        return self._save_dir
+
+    def __deepcopy__(self, memo=None):
+        """Override of the deepcopy method."""
+        return self

eva/core/loggers/experimental_loggers.py

@@ -0,0 +1,8 @@
+"""Experiment loggers."""
+
+from typing import Union
+
+from lightning.pytorch.loggers import CSVLogger, TensorBoardLogger
+
+"""Supported loggers."""
+ExperimentalLoggers = Union[CSVLogger, TensorBoardLogger]

eva/core/loggers/log/__init__.py

@@ -0,0 +1,5 @@
+"""Experiment loggers actions."""
+
+from eva.core.loggers.log.parameters import log_parameters
+
+__all__ = ["log_parameters"]