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

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

@@ -1,10 +1,10 @@
 """Base dataset class for Embeddings."""
 
 import abc
+import multiprocessing
 import os
-from typing import Callable, Dict, Literal, Tuple
+from typing import Callable, Dict, Generic, Literal, Tuple, TypeVar
 
-import numpy as np
 import pandas as pd
 import torch
 from typing_extensions import override
@@ -12,16 +12,20 @@ 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": "slide_id",
+    "multi_id": "wsi_id",
 }
 """The default column mapping of the variables to the manifest columns."""
 
 
-class EmbeddingsDataset(base.Dataset):
+class EmbeddingsDataset(base.Dataset, Generic[TargetType]):
     """Abstract base class for embedding datasets."""
 
     def __init__(
@@ -62,31 +66,7 @@ class EmbeddingsDataset(base.Dataset):
 
         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."""
+        self._set_multiprocessing_start_method()
 
     def filename(self, index: int) -> str:
         """Returns the filename of the `index`'th data sample.
@@ -105,7 +85,11 @@ class EmbeddingsDataset(base.Dataset):
     def setup(self):
         self._data = self._load_manifest()
 
-    def __getitem__(self, index) -> Tuple[torch.Tensor, np.ndarray]:
+    @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:
@@ -118,6 +102,28 @@ class EmbeddingsDataset(base.Dataset):
         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.
 
@@ -132,8 +138,8 @@ class EmbeddingsDataset(base.Dataset):
         return data
 
     def _apply_transforms(
-        self, embeddings: torch.Tensor, target: np.ndarray
-    ) -> Tuple[torch.Tensor, np.ndarray]:
+        self, embeddings: torch.Tensor, target: TargetType
+    ) -> Tuple[torch.Tensor, TargetType]:
         """Applies the transforms to the provided data and returns them.
 
         Args:
@@ -150,3 +156,12 @@ class EmbeddingsDataset(base.Dataset):
             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/loggers/experimental_loggers.py

@@ -2,7 +2,7 @@
 
 from typing import Union
 
-from lightning.pytorch.loggers import CSVLogger, TensorBoardLogger
+from lightning.pytorch.loggers import CSVLogger, TensorBoardLogger, WandbLogger
 
 """Supported loggers."""
-ExperimentalLoggers = Union[CSVLogger, TensorBoardLogger]
+ExperimentalLoggers = Union[CSVLogger, TensorBoardLogger, WandbLogger]

eva/core/loggers/log/__init__.py

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

eva/core/loggers/log/image.py

@@ -0,0 +1,71 @@
+"""Image log functionality."""
+
+import functools
+
+import torch
+
+from eva.core.loggers import loggers
+from eva.core.loggers.log import utils
+
+
+@functools.singledispatch
+def log_image(
+    logger,
+    tag: str,
+    image: torch.Tensor,
+    step: int = 0,
+) -> None:
+    """Adds an image to the logger.
+
+    Args:
+        logger: The desired logger.
+        tag: The log tag.
+        image: The image tensor to log. It should have
+            the shape of (3,H,W) and (0,1) normalized.
+        step: The global step of the log.
+    """
+    utils.raise_not_supported(logger, "image")
+
+
+@log_image.register
+def _(
+    loggers: list,
+    tag: str,
+    image: torch.Tensor,
+    step: int = 0,
+) -> None:
+    """Adds an image to a list of supported loggers."""
+    for logger in loggers:
+        log_image(
+            logger,
+            tag=tag,
+            image=image,
+            step=step,
+        )
+
+
+@log_image.register
+def _(
+    logger: loggers.TensorBoardLogger,
+    tag: str,
+    image: torch.Tensor,
+    step: int = 0,
+) -> None:
+    """Adds an image to a TensorBoard logger."""
+    logger.experiment.add_image(
+        tag=tag,
+        img_tensor=image,
+        global_step=step,
+    )
+
+
+@log_image.register
+def _(
+    logger: loggers.WandbLogger,
+    tag: str,
+    image: torch.Tensor,
+    caption: str | None = None,
+    step: int = 0,
+) -> None:
+    """Adds a list of images to a Wandb logger."""
+    logger.log_image(key=tag, images=[image.float()], step=step, caption=[caption])

eva/core/loggers/log/parameters.py

@@ -51,6 +51,16 @@ def _(
     )
 
 
+@log_parameters.register
+def _(
+    logger: loggers_lib.WandbLogger,
+    tag: str,
+    parameters: Dict[str, Any],
+) -> None:
+    """Adds parameters to a Wandb logger."""
+    logger.experiment.config.update(parameters)
+
+
 def _yaml_to_markdown(data: Dict[str, Any]) -> str:
     """Casts yaml data to markdown.
 

eva/core/loggers/loggers.py

@@ -0,0 +1,6 @@
+"""Experimental loggers."""
+
+from lightning.pytorch.loggers import TensorBoardLogger, WandbLogger
+
+Loggers = TensorBoardLogger | WandbLogger
+"""Supported loggers."""

eva/core/metrics/__init__.py

@@ -3,15 +3,19 @@
 from eva.core.metrics.average_loss import AverageLoss
 from eva.core.metrics.binary_balanced_accuracy import BinaryBalancedAccuracy
 from eva.core.metrics.defaults import BinaryClassificationMetrics, MulticlassClassificationMetrics
+from eva.core.metrics.generalized_dice import GeneralizedDiceScore
+from eva.core.metrics.mean_iou import MeanIoU
 from eva.core.metrics.structs import Metric, MetricCollection, MetricModule, MetricsSchema
 
 __all__ = [
     "AverageLoss",
     "BinaryBalancedAccuracy",
+    "BinaryClassificationMetrics",
+    "MulticlassClassificationMetrics",
+    "GeneralizedDiceScore",
+    "MeanIoU",
     "Metric",
     "MetricCollection",
     "MetricModule",
     "MetricsSchema",
-    "MulticlassClassificationMetrics",
-    "BinaryClassificationMetrics",
 ]

eva/core/metrics/defaults/__init__.py

@@ -1,6 +1,13 @@
 """Default metric collections API."""
 
-from eva.core.metrics.defaults.classification.binary import BinaryClassificationMetrics
-from eva.core.metrics.defaults.classification.multiclass import MulticlassClassificationMetrics
+from eva.core.metrics.defaults.classification import (
+    BinaryClassificationMetrics,
+    MulticlassClassificationMetrics,
+)
+from eva.core.metrics.defaults.segmentation import MulticlassSegmentationMetrics
 
-__all__ = ["MulticlassClassificationMetrics", "BinaryClassificationMetrics"]
+__all__ = [
+    "MulticlassClassificationMetrics",
+    "BinaryClassificationMetrics",
+    "MulticlassSegmentationMetrics",
+]

eva/core/metrics/defaults/classification/__init__.py

@@ -3,4 +3,4 @@
 from eva.core.metrics.defaults.classification.binary import BinaryClassificationMetrics
 from eva.core.metrics.defaults.classification.multiclass import MulticlassClassificationMetrics
 
-__all__ = ["MulticlassClassificationMetrics", "BinaryClassificationMetrics"]
+__all__ = ["BinaryClassificationMetrics", "MulticlassClassificationMetrics"]

eva/core/metrics/defaults/classification/binary.py

@@ -17,15 +17,6 @@ class BinaryClassificationMetrics(structs.MetricCollection):
     ) -> None:
         """Initializes the binary classification metrics.
 
-        The metrics instantiated here are:
-
-        - BinaryAUROC
-        - BinaryAccuracy
-        - BinaryBalancedAccuracy
-        - BinaryF1Score
-        - BinaryPrecision
-        - BinaryRecall
-
         Args:
             threshold: Threshold for transforming probability to binary (0,1) predictions
             ignore_index: Specifies a target value that is ignored and does not

eva/core/metrics/defaults/classification/multiclass.py

@@ -20,14 +20,6 @@ class MulticlassClassificationMetrics(structs.MetricCollection):
     ) -> None:
         """Initializes the multi-class classification metrics.
 
-        The metrics instantiated here are:
-
-        - MulticlassAccuracy
-        - MulticlassPrecision
-        - MulticlassRecall
-        - MulticlassF1Score
-        - MulticlassAUROC
-
         Args:
             num_classes: Integer specifying the number of classes.
             average: Defines the reduction that is applied over labels.

eva/core/metrics/defaults/segmentation/__init__.py

@@ -0,0 +1,5 @@
+"""Default segmentation metric collections API."""
+
+from eva.core.metrics.defaults.segmentation.multiclass import MulticlassSegmentationMetrics
+
+__all__ = ["MulticlassSegmentationMetrics"]

eva/core/metrics/defaults/segmentation/multiclass.py

@@ -0,0 +1,43 @@
+"""Default metric collection for multiclass semantic segmentation tasks."""
+
+from eva.core.metrics import generalized_dice, mean_iou, structs
+
+
+class MulticlassSegmentationMetrics(structs.MetricCollection):
+    """Default metrics for multi-class semantic segmentation tasks."""
+
+    def __init__(
+        self,
+        num_classes: int,
+        include_background: bool = False,
+        ignore_index: int | None = None,
+        prefix: str | None = None,
+        postfix: str | None = None,
+    ) -> None:
+        """Initializes the multi-class semantic segmentation metrics.
+
+        Args:
+            num_classes: Integer specifying the number of classes.
+            include_background: Whether to include the background class in the metrics computation.
+            ignore_index: Integer specifying a target class to ignore. If given, this class
+                index does not contribute to the returned score, regardless of reduction method.
+            prefix: A string to add before the keys in the output dictionary.
+            postfix: A string to add after the keys in the output dictionary.
+        """
+        super().__init__(
+            metrics=[
+                generalized_dice.GeneralizedDiceScore(
+                    num_classes=num_classes,
+                    include_background=include_background,
+                    weight_type="linear",
+                    ignore_index=ignore_index,
+                ),
+                mean_iou.MeanIoU(
+                    num_classes=num_classes,
+                    include_background=include_background,
+                    ignore_index=ignore_index,
+                ),
+            ],
+            prefix=prefix,
+            postfix=postfix,
+        )

eva/core/metrics/generalized_dice.py

@@ -0,0 +1,59 @@
+"""Generalized Dice Score metric for semantic segmentation."""
+
+from typing import Any, Literal
+
+import torch
+from torchmetrics import segmentation
+from typing_extensions import override
+
+
+class GeneralizedDiceScore(segmentation.GeneralizedDiceScore):
+    """Defines the Generalized Dice Score.
+
+    It expands the `torchmetrics` class by including an `ignore_index`
+    functionality.
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        include_background: bool = True,
+        weight_type: Literal["square", "simple", "linear"] = "linear",
+        ignore_index: int | None = None,
+        per_class: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """Initializes the metric.
+
+        Args:
+            num_classes: The number of classes in the segmentation problem.
+            include_background: Whether to include the background class in the computation
+            weight_type: The type of weight to apply to each class. Can be one of `"square"`,
+                `"simple"`, or `"linear"`.
+            input_format: What kind of input the function receives. Choose between ``"one-hot"``
+                for one-hot encoded tensors or ``"index"`` for index tensors.
+            ignore_index: Integer specifying a target class to ignore. If given, this class
+                index does not contribute to the returned score, regardless of reduction method.
+            per_class: Whether to compute the IoU for each class separately. If set to ``False``,
+                the metric will compute the mean IoU over all classes.
+            kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
+        """
+        super().__init__(
+            num_classes=num_classes,
+            include_background=include_background,
+            weight_type=weight_type,
+            per_class=per_class,
+            **kwargs,
+        )
+
+        self.ignore_index = ignore_index
+
+    @override
+    def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
+        if self.ignore_index is not None:
+            mask = target != self.ignore_index
+            mask = mask.all(dim=-1, keepdim=True)
+            preds = preds * mask
+            target = target * mask
+
+        super().update(preds=preds, target=target)

eva/core/metrics/mean_iou.py

@@ -0,0 +1,120 @@
+"""Mean Intersection over Union (mIoU) metric for semantic segmentation."""
+
+from typing import Any, Literal, Tuple
+
+import torch
+import torchmetrics
+
+
+class MeanIoU(torchmetrics.Metric):
+    """Computes Mean Intersection over Union (mIoU) for semantic segmentation.
+
+    Fixes the torchmetrics implementation
+    (issue https://github.com/Lightning-AI/torchmetrics/issues/2558)
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        include_background: bool = True,
+        ignore_index: int | None = None,
+        per_class: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """Initializes the metric.
+
+        Args:
+            num_classes: The number of classes in the segmentation problem.
+            include_background: Whether to include the background class in the computation
+            ignore_index: Integer specifying a target class to ignore. If given, this class
+                index does not contribute to the returned score, regardless of reduction method.
+            per_class: Whether to compute the IoU for each class separately. If set to ``False``,
+                the metric will compute the mean IoU over all classes.
+            kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
+        """
+        super().__init__(**kwargs)
+
+        self.num_classes = num_classes
+        self.include_background = include_background
+        self.ignore_index = ignore_index
+        self.per_class = per_class
+
+        self.add_state("intersection", default=torch.zeros(num_classes), dist_reduce_fx="sum")
+        self.add_state("union", default=torch.zeros(num_classes), dist_reduce_fx="sum")
+
+    def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
+        """Update the state with the new data."""
+        intersection, union = _compute_intersection_and_union(
+            preds,
+            target,
+            num_classes=self.num_classes,
+            include_background=self.include_background,
+            ignore_index=self.ignore_index,
+        )
+        self.intersection += intersection.sum(0)
+        self.union += union.sum(0)
+
+    def compute(self) -> torch.Tensor:
+        """Compute the final mean IoU score."""
+        iou_valid = torch.gt(self.union, 0)
+        iou = torch.where(
+            iou_valid,
+            torch.divide(self.intersection, self.union),
+            torch.nan,
+        )
+        if not self.per_class:
+            iou = torch.mean(iou[iou_valid])
+        return iou
+
+
+def _compute_intersection_and_union(
+    preds: torch.Tensor,
+    target: torch.Tensor,
+    num_classes: int,
+    include_background: bool = False,
+    input_format: Literal["one-hot", "index"] = "index",
+    ignore_index: int | None = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Compute the intersection and union for semantic segmentation tasks.
+
+    Args:
+        preds: Predicted tensor with shape (N, ...) where N is the batch size.
+            The shape can be (N, H, W) for 2D data or (N, D, H, W) for 3D data.
+        target: Ground truth tensor with the same shape as preds.
+        num_classes: Number of classes in the segmentation task.
+        include_background: Whether to include the background class in the computation.
+        input_format: Format of the input tensors.
+        ignore_index: Integer specifying a target class to ignore. If given, this class
+            index does not contribute to the returned score, regardless of reduction method.
+
+    Returns:
+        Two tensors representing the intersection and union for each class.
+        Shape of each tensor is (N, num_classes).
+
+    Note:
+        - If input_format is "index", the tensors are converted to one-hot encoding.
+        - If include_background is `False`, the background class
+          (assumed to be the first channel) is ignored in the computation.
+    """
+    if ignore_index is not None:
+        mask = target != ignore_index
+        mask = mask.all(dim=-1, keepdim=True)
+        preds = preds * mask
+        target = target * mask
+
+    if input_format == "index":
+        preds = torch.nn.functional.one_hot(preds, num_classes=num_classes)
+        target = torch.nn.functional.one_hot(target, num_classes=num_classes)
+
+    if not include_background:
+        preds[..., 0] = 0
+        target[..., 0] = 0
+
+    reduce_axis = list(range(1, preds.ndim - 1))
+
+    intersection = torch.sum(torch.logical_and(preds, target), dim=reduce_axis)
+    target_sum = torch.sum(target, dim=reduce_axis)
+    pred_sum = torch.sum(preds, dim=reduce_axis)
+    union = target_sum + pred_sum - intersection
+
+    return intersection, union

eva/core/metrics/structs/schemas.py

@@ -44,4 +44,6 @@ class MetricsSchema:
         if metrics is None or self.common is None:
             return self.common or metrics
 
-        return [self.common, metrics]  # type: ignore
+        metrics = metrics if isinstance(metrics, list) else [metrics]  # type: ignore
+        common = self.common if isinstance(self.common, list) else [self.common]
+        return common + metrics  # type: ignore