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

eva/core/callbacks/__init__.py

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

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/callbacks/writers/__init__.py

@@ -1,5 +1,8 @@
-"""Callbacks API."""
+"""Writers callbacks API."""
 
-from eva.core.callbacks.writers.embeddings import EmbeddingsWriter
+from eva.core.callbacks.writers.embeddings import (
+    ClassificationEmbeddingsWriter,
+    SegmentationEmbeddingsWriter,
+)
 
-__all__ = ["EmbeddingsWriter"]
+__all__ = ["ClassificationEmbeddingsWriter", "SegmentationEmbeddingsWriter"]

eva/core/callbacks/writers/embeddings/__init__.py

@@ -0,0 +1,6 @@
+"""Embedding callback writers."""
+
+from eva.core.callbacks.writers.embeddings.classification import ClassificationEmbeddingsWriter
+from eva.core.callbacks.writers.embeddings.segmentation import SegmentationEmbeddingsWriter
+
+__all__ = ["ClassificationEmbeddingsWriter", "SegmentationEmbeddingsWriter"]

eva/core/callbacks/writers/embeddings/_manifest.py

@@ -0,0 +1,71 @@
+"""Manifest file manager."""
+
+import csv
+import io
+import os
+from typing import Any, Dict, List
+
+import _csv
+import torch
+
+
+class ManifestManager:
+    """Class for writing the embedding manifest files."""
+
+    def __init__(
+        self,
+        output_dir: str,
+        metadata_keys: List[str] | None = None,
+        overwrite: bool = False,
+    ) -> None:
+        """Initializes the writing manager.
+
+        Args:
+            output_dir: The directory where the embeddings will be saved.
+            metadata_keys: An optional list of keys to extract from the batch
+                metadata and store as additional columns in the manifest file.
+            overwrite: Whether to overwrite the output directory.
+        """
+        self._output_dir = output_dir
+        self._metadata_keys = metadata_keys or []
+        self._overwrite = overwrite
+
+        self._manifest_file: io.TextIOWrapper
+        self._manifest_writer: _csv.Writer  # type: ignore
+
+        self._setup()
+
+    def _setup(self) -> None:
+        """Initializes the manifest file and sets the file object and writer."""
+        manifest_path = os.path.join(self._output_dir, "manifest.csv")
+        if os.path.exists(manifest_path) and not self._overwrite:
+            raise FileExistsError(
+                f"A manifest file already exists at {manifest_path}, which indicates that the "
+                "chosen output directory has been previously used for writing embeddings."
+            )
+        self._manifest_file = open(manifest_path, "w", newline="")
+        self._manifest_writer = csv.writer(self._manifest_file)
+        self._manifest_writer.writerow(
+            ["origin", "embeddings", "target", "split"] + self._metadata_keys
+        )
+
+    def update(
+        self,
+        input_name: str,
+        save_name: str,
+        target: str,
+        split: str | None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> None:
+        """Adds a new entry to the manifest file."""
+        metadata_entries = _to_dict_values(metadata or {})
+        self._manifest_writer.writerow([input_name, save_name, target, split] + metadata_entries)
+
+    def close(self) -> None:
+        """Closes the manifest file."""
+        if self._manifest_file:
+            self._manifest_file.close()
+
+
+def _to_dict_values(data: Dict[str, Any]) -> List[Any]:
+    return [value.item() if isinstance(value, torch.Tensor) else value for value in data.values()]

eva/core/callbacks/writers/embeddings/base.py

@@ -0,0 +1,192 @@
+"""Embeddings writer base class."""
+
+import abc
+import io
+import os
+from typing import Any, Dict, List, Sequence
+
+import lightning.pytorch as pl
+import torch
+from lightning.pytorch import callbacks
+from loguru import logger
+from torch import multiprocessing, nn
+from typing_extensions import override
+
+from eva.core import utils
+from eva.core.callbacks.writers.embeddings.typings import QUEUE_ITEM
+from eva.core.models.modules.typings import INPUT_BATCH
+from eva.core.utils import multiprocessing as eva_multiprocessing
+
+
+class EmbeddingsWriter(callbacks.BasePredictionWriter, abc.ABC):
+    """Callback for writing generated embeddings to disk."""
+
+    def __init__(
+        self,
+        output_dir: str,
+        backbone: nn.Module | None = None,
+        dataloader_idx_map: Dict[int, str] | None = None,
+        metadata_keys: List[str] | None = None,
+        overwrite: bool = False,
+        save_every_n: int = 100,
+    ) -> None:
+        """Initializes a new EmbeddingsWriter instance.
+
+        This callback writes the embedding files in a separate process to avoid blocking the
+        main process where the model forward pass is executed.
+
+        Args:
+            output_dir: The directory where the embeddings will be saved.
+            backbone: A model to be used as feature extractor. If `None`,
+                it will be expected that the input batch returns the features directly.
+            dataloader_idx_map: A dictionary mapping dataloader indices to their respective
+                names (e.g. train, val, test).
+            metadata_keys: An optional list of keys to extract from the batch metadata and store
+                as additional columns in the manifest file.
+            overwrite: Whether to overwrite if embeddings are already present in the specified
+                output directory. If set to `False`, an error will be raised if embeddings are
+                already present (recommended).
+            save_every_n: Interval for number of iterations to save the embeddings to disk.
+                During this interval, the embeddings are accumulated in memory.
+        """
+        super().__init__(write_interval="batch")
+
+        self._output_dir = output_dir
+        self._backbone = backbone
+        self._dataloader_idx_map = dataloader_idx_map or {}
+        self._overwrite = overwrite
+        self._save_every_n = save_every_n
+        self._metadata_keys = metadata_keys or []
+
+        self._write_queue: multiprocessing.Queue
+        self._write_process: eva_multiprocessing.Process
+
+    @staticmethod
+    @abc.abstractmethod
+    def _process_write_queue(
+        write_queue: multiprocessing.Queue,
+        output_dir: str,
+        metadata_keys: List[str],
+        save_every_n: int,
+        overwrite: bool = False,
+    ) -> None:
+        """This function receives and processes items added by the main process to the queue.
+
+        Queue items contain the embedding tensors, targets and metadata which need to be
+        saved to disk (.pt files and manifest).
+        """
+
+    @override
+    def on_predict_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        self._check_if_exists()
+        self._initialize_write_process()
+        self._write_process.start()
+
+        if self._backbone is not None:
+            self._backbone = self._backbone.to(pl_module.device)
+            self._backbone.eval()
+
+    @override
+    def write_on_batch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        prediction: Any,
+        batch_indices: Sequence[int],
+        batch: INPUT_BATCH,
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+        dataset = trainer.predict_dataloaders[dataloader_idx].dataset  # type: ignore
+        _, targets, metadata = INPUT_BATCH(*batch)
+        split = self._dataloader_idx_map.get(dataloader_idx)
+        if not isinstance(targets, torch.Tensor):
+            raise ValueError(f"Targets ({type(targets)}) should be `torch.Tensor`.")
+
+        with torch.no_grad():
+            embeddings = self._get_embeddings(prediction)
+
+        for local_idx, global_idx in enumerate(batch_indices[: len(embeddings)]):
+            data_name = dataset.filename(global_idx)
+            save_name = os.path.splitext(data_name)[0] + ".pt"
+            embeddings_buffer, target_buffer = _as_io_buffers(
+                embeddings[local_idx], targets[local_idx]
+            )
+            item_metadata = self._get_item_metadata(metadata, local_idx)
+            item = QUEUE_ITEM(
+                prediction_buffer=embeddings_buffer,
+                target_buffer=target_buffer,
+                data_name=data_name,
+                save_name=save_name,
+                split=split,
+                metadata=item_metadata,
+            )
+            self._write_queue.put(item)
+
+        self._write_process.check_exceptions()
+
+    @override
+    def on_predict_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        self._write_queue.put(None)
+        self._write_process.join()
+        logger.info(f"Predictions and manifest saved to {self._output_dir}")
+
+    def _initialize_write_process(self) -> None:
+        self._write_queue = multiprocessing.Queue()
+        self._write_process = eva_multiprocessing.Process(
+            target=self._process_write_queue,
+            args=(
+                self._write_queue,
+                self._output_dir,
+                self._metadata_keys,
+                self._save_every_n,
+                self._overwrite,
+            ),
+        )
+
+    @abc.abstractmethod
+    def _get_embeddings(self, tensor: torch.Tensor) -> torch.Tensor | List[List[torch.Tensor]]:
+        """Returns the embeddings from predictions."""
+
+    def _get_item_metadata(
+        self, metadata: Dict[str, Any] | None, local_idx: int
+    ) -> Dict[str, Any] | None:
+        """Returns the metadata for the item at the given local index."""
+        if not metadata:
+            if self._metadata_keys:
+                raise ValueError("Metadata keys are provided but the batch metadata is empty.")
+            else:
+                return None
+
+        item_metadata = {}
+        for key in self._metadata_keys:
+            if key not in metadata:
+                raise KeyError(f"Metadata key '{key}' not found in the batch metadata.")
+            metadata_value = metadata[key][local_idx]
+            try:
+                item_metadata[key] = utils.to_cpu(metadata_value)
+            except TypeError:
+                item_metadata[key] = metadata_value
+
+        return item_metadata
+
+    def _check_if_exists(self) -> None:
+        """Checks if the output directory already exists and if it should be overwritten."""
+        try:
+            os.makedirs(self._output_dir, exist_ok=self._overwrite)
+        except FileExistsError as e:
+            raise FileExistsError(
+                f"The embeddings output directory already exists: {self._output_dir}. This "
+                "either means that they have been computed before or that a wrong output "
+                "directory is being used. Consider using `eva fit` instead, selecting a "
+                "different output directory or setting overwrite=True."
+            ) from e
+        os.makedirs(self._output_dir, exist_ok=True)
+
+
+def _as_io_buffers(*items: torch.Tensor | List[torch.Tensor]) -> Sequence[io.BytesIO]:
+    """Loads torch tensors as io buffers."""
+    buffers = [io.BytesIO() for _ in range(len(items))]
+    for tensor, buffer in zip(items, buffers, strict=False):
+        torch.save(utils.clone(tensor), buffer)
+    return buffers

eva/core/callbacks/writers/embeddings/classification.py

@@ -0,0 +1,117 @@
+"""Embeddings writer for classification."""
+
+import io
+import os
+from typing import Dict, List
+
+import torch
+from torch import multiprocessing
+from typing_extensions import override
+
+from eva.core.callbacks.writers.embeddings import base
+from eva.core.callbacks.writers.embeddings._manifest import ManifestManager
+from eva.core.callbacks.writers.embeddings.typings import ITEM_DICT_ENTRY, QUEUE_ITEM
+
+
+class ClassificationEmbeddingsWriter(base.EmbeddingsWriter):
+    """Callback for writing generated embeddings to disk for classification tasks."""
+
+    @staticmethod
+    @override
+    def _process_write_queue(
+        write_queue: multiprocessing.Queue,
+        output_dir: str,
+        metadata_keys: List[str],
+        save_every_n: int,
+        overwrite: bool = False,
+    ) -> None:
+        """Processes the write queue and saves the predictions to disk.
+
+        Note that in Multi Instance Learning (MIL) scenarios, we can have multiple
+        embeddings per input data point. In that case, this function will save all
+        embeddings that correspond to the same data point as a list of tensors to
+        the same .pt file.
+        """
+        manifest_manager = ManifestManager(output_dir, metadata_keys, overwrite)
+        name_to_items: Dict[str, ITEM_DICT_ENTRY] = {}
+
+        counter = 0
+        while True:
+            item = write_queue.get()
+            if item is None:
+                break
+            item = QUEUE_ITEM(*item)
+
+            if item.save_name in name_to_items:
+                name_to_items[item.save_name].items.append(item)
+            else:
+                name_to_items[item.save_name] = ITEM_DICT_ENTRY(items=[item], save_count=0)
+
+            if counter > 0 and counter % save_every_n == 0:
+                name_to_items = _save_items(name_to_items, output_dir, manifest_manager)
+            counter += 1
+
+        if len(name_to_items) > 0:
+            _save_items(name_to_items, output_dir, manifest_manager)
+
+        manifest_manager.close()
+
+    @override
+    def _get_embeddings(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Returns the embeddings from predictions."""
+        return self._backbone(tensor) if self._backbone else tensor
+
+
+def _save_items(
+    name_to_items: Dict[str, ITEM_DICT_ENTRY],
+    output_dir: str,
+    manifest_manager: ManifestManager,
+) -> Dict[str, ITEM_DICT_ENTRY]:
+    """Saves predictions to disk and updates the manifest file.
+
+    Args:
+        name_to_items: A dictionary mapping save data point names to the corresponding queue items
+            holding the prediction tensors and the information for the manifest file.
+        output_dir: The directory where the embedding tensors & manifest will be saved.
+        manifest_manager: The manifest manager instance to update the manifest file.
+    """
+    for save_name, entry in name_to_items.items():
+        if len(entry.items) > 0:
+            save_path = os.path.join(output_dir, save_name)
+            is_first_save = entry.save_count == 0
+            if is_first_save:
+                _, target, input_name, _, split, metadata = QUEUE_ITEM(*entry.items[0])
+                target = torch.load(io.BytesIO(target.getbuffer()), map_location="cpu").item()
+                manifest_manager.update(input_name, save_name, target, split, metadata)
+
+            prediction_buffers = [item.prediction_buffer for item in entry.items]
+            _save_predictions(prediction_buffers, save_path, is_first_save)
+            name_to_items[save_name].save_count += 1
+            name_to_items[save_name].items = []
+
+    return name_to_items
+
+
+def _save_predictions(
+    prediction_buffers: List[io.BytesIO], save_path: str, is_first_save: bool
+) -> None:
+    """Saves the embedding tensors as list to .pt files.
+
+    If it's not the first save to this save_path, the new predictions are appended to
+    the existing ones and saved to the same file.
+
+    Example use-case: Save all patch embeddings corresponding to the same WSI to a single file.
+    """
+    predictions = [
+        torch.load(io.BytesIO(buffer.getbuffer()), map_location="cpu")
+        for buffer in prediction_buffers
+    ]
+
+    if not is_first_save:
+        previous_predictions = torch.load(save_path, map_location="cpu")
+        if not isinstance(previous_predictions, list):
+            raise ValueError("Previous predictions should be a list of tensors.")
+        predictions = predictions + previous_predictions
+
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    torch.save(predictions, save_path)

eva/core/callbacks/writers/embeddings/segmentation.py

@@ -0,0 +1,78 @@
+"""Segmentation embeddings writer."""
+
+import collections
+import io
+import os
+from typing import List
+
+import torch
+from torch import multiprocessing
+from typing_extensions import override
+
+from eva.core.callbacks.writers.embeddings import base
+from eva.core.callbacks.writers.embeddings._manifest import ManifestManager
+from eva.core.callbacks.writers.embeddings.typings import QUEUE_ITEM
+
+
+class SegmentationEmbeddingsWriter(base.EmbeddingsWriter):
+    """Callback for writing generated embeddings to disk."""
+
+    @staticmethod
+    @override
+    def _process_write_queue(
+        write_queue: multiprocessing.Queue,
+        output_dir: str,
+        metadata_keys: List[str],
+        save_every_n: int,
+        overwrite: bool = False,
+    ) -> None:
+        manifest_manager = ManifestManager(output_dir, metadata_keys, overwrite)
+        counter = collections.defaultdict(lambda: -1)
+        while True:
+            item = write_queue.get()
+            if item is None:
+                break
+
+            embeddings_buffer, target_buffer, input_name, save_name, split, metadata = QUEUE_ITEM(
+                *item
+            )
+            counter[save_name] += 1
+            save_name = save_name.replace(".pt", f"-{counter[save_name]}.pt")
+            target_filename = save_name.replace(".pt", "-mask.pt")
+
+            _save_embedding(embeddings_buffer, save_name, output_dir)
+            _save_embedding(target_buffer, target_filename, output_dir)
+            manifest_manager.update(input_name, save_name, target_filename, split, metadata)
+
+        manifest_manager.close()
+
+    @override
+    def _get_embeddings(self, tensor: torch.Tensor) -> torch.Tensor | List[List[torch.Tensor]]:
+        """Returns the embeddings from predictions."""
+
+        def _get_grouped_embeddings(embeddings: List[torch.Tensor]) -> List[List[torch.Tensor]]:
+            """Casts a list of multi-leveled batched embeddings to grouped per batch.
+
+            That is, for embeddings to be a list of shape (batch_size, hidden_dim, height, width),
+            such as `[(2, 192, 16, 16), (2, 192, 16, 16)]`, to be reshaped as a list of lists of
+            per batch multi-level embeddings, thus
+            `[ [(192, 16, 16), (192, 16, 16)], [(192, 16, 16), (192, 16, 16)] ]`.
+            """
+            batch_size = embeddings[0].shape[0]
+            grouped_embeddings = []
+            for batch_idx in range(batch_size):
+                batch_list = [layer_embeddings[batch_idx] for layer_embeddings in embeddings]
+                grouped_embeddings.append(batch_list)
+            return grouped_embeddings
+
+        embeddings = self._backbone(tensor) if self._backbone else tensor
+        if isinstance(embeddings, list):
+            embeddings = _get_grouped_embeddings(embeddings)
+        return embeddings
+
+
+def _save_embedding(embeddings_buffer: io.BytesIO, save_name: str, output_dir: str) -> None:
+    save_path = os.path.join(output_dir, save_name)
+    prediction = torch.load(io.BytesIO(embeddings_buffer.getbuffer()), map_location="cpu")
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    torch.save(prediction, save_path)

eva/core/callbacks/writers/embeddings/typings.py

@@ -0,0 +1,38 @@
+"""Typing definitions for the writer callback functions."""
+
+import dataclasses
+import io
+from typing import Any, Dict, List, NamedTuple
+
+
+class QUEUE_ITEM(NamedTuple):
+    """The default input batch data scheme."""
+
+    prediction_buffer: io.BytesIO
+    """IO buffer containing the prediction tensor."""
+
+    target_buffer: io.BytesIO
+    """IO buffer containing the target tensor."""
+
+    data_name: str
+    """Name of the input data that was used to generate the embedding."""
+
+    save_name: str
+    """Name to store the generated embedding."""
+
+    split: str | None
+    """The dataset split the item belongs to (e.g. train, val, test)."""
+
+    metadata: Dict[str, Any] | None = None
+    """Dictionary holding additional metadata."""
+
+
+@dataclasses.dataclass
+class ITEM_DICT_ENTRY:
+    """Typing for holding queue items and number of save operations."""
+
+    items: List[QUEUE_ITEM]
+    """List of queue items."""
+
+    save_count: int
+    """Number of prior item batch saves to same file."""

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.classification import (
+    EmbeddingsClassificationDataset,
+    MultiEmbeddingsClassificationDataset,
+)
 from eva.core.data.datasets.dataset import TorchDataset
 
-__all__ = ["Dataset", "EmbeddingsClassificationDataset", "TorchDataset"]
+__all__ = [
+    "Dataset",
+    "EmbeddingsClassificationDataset",
+    "MultiEmbeddingsClassificationDataset",
+    "TorchDataset",
+]