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

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, WandbLogger
+
+"""Supported loggers."""
+ExperimentalLoggers = Union[CSVLogger, TensorBoardLogger, WandbLogger]

eva/core/loggers/log/__init__.py

@@ -0,0 +1,6 @@
+"""Experiment loggers operations."""
+
+from eva.core.loggers.log.image import log_image
+from eva.core.loggers.log.parameters import 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

@@ -0,0 +1,74 @@
+"""Text log functionality."""
+
+import functools
+from typing import Any, Dict
+
+import yaml
+
+from eva.core.loggers import experimental_loggers as loggers_lib
+from eva.core.loggers.log import utils
+
+
+@functools.singledispatch
+def log_parameters(
+    logger,
+    tag: str,
+    parameters: Dict[str, Any],
+) -> None:
+    """Adds parameters to the logger.
+
+    Args:
+        logger: The desired logger.
+        tag: The log tag.
+        parameters: The parameters to log.
+    """
+    utils.raise_not_supported(logger, "parameters")
+
+
+@log_parameters.register
+def _(
+    loggers: list,
+    tag: str,
+    parameters: Dict[str, Any],
+) -> None:
+    """Adds parameters to a list of supported loggers."""
+    for logger in loggers:
+        log_parameters(logger, tag=tag, parameters=parameters)
+
+
+@log_parameters.register
+def _(
+    logger: loggers_lib.TensorBoardLogger,
+    tag: str,
+    parameters: Dict[str, Any],
+) -> None:
+    """Adds parameters to a TensorBoard logger."""
+    as_markdown_text = _yaml_to_markdown(parameters)
+    logger.experiment.add_text(
+        tag=tag,
+        text_string=as_markdown_text,
+        global_step=0,
+    )
+
+
+@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.
+
+    Args:
+        data: The yaml data.
+
+    Returns:
+        A string markdown friendly formatted.
+    """
+    text = yaml.dump(data, sort_keys=False)
+    return f"```yaml\n{text}```"

eva/core/loggers/log/utils.py

@@ -0,0 +1,13 @@
+"""Logging related utilities."""
+
+from loguru import logger as cli_logger
+
+from eva.core.loggers import ExperimentalLoggers
+
+
+def raise_not_supported(logger: ExperimentalLoggers, data_type: str) -> None:
+    """Raises a warning for not supported tasks from the given logger."""
+    print("\n")
+    cli_logger.debug(
+        f"Logger '{logger.__class__.__name__}' is not supported for " f"'{data_type}' data."
+    )

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

eva/core/models/__init__.py

@@ -1,12 +1,14 @@
 """Models API."""
 
 from eva.core.models.modules import HeadModule, InferenceModule
-from eva.core.models.networks import MLP, HuggingFaceModel, ModelFromFunction, ONNXModel
+from eva.core.models.networks import MLP
+from eva.core.models.wrappers import BaseModel, HuggingFaceModel, ModelFromFunction, ONNXModel
 
 __all__ = [
     "HeadModule",
     "InferenceModule",
     "MLP",
+    "BaseModel",
     "HuggingFaceModel",
     "ModelFromFunction",
     "ONNXModel",

eva/core/models/modules/head.py

@@ -1,11 +1,11 @@
 """"Neural Network Head Module."""
 
-from typing import Any, Callable
+from typing import Any, Callable, Dict
 
 import torch
 from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
 from lightning.pytorch.utilities.types import STEP_OUTPUT
-from torch import optim
+from torch import nn, optim
 from torch.optim import lr_scheduler
 from typing_extensions import override
 
@@ -13,6 +13,7 @@ from eva.core.metrics import structs as metrics_lib
 from eva.core.models.modules import module
 from eva.core.models.modules.typings import INPUT_BATCH, MODEL_TYPE
 from eva.core.models.modules.utils import batch_postprocess, grad
+from eva.core.utils import parser
 
 
 class HeadModule(module.ModelModule):
@@ -24,7 +25,7 @@ class HeadModule(module.ModelModule):
 
     def __init__(
         self,
-        head: MODEL_TYPE,
+        head: Dict[str, Any] | MODEL_TYPE,
         criterion: Callable[..., torch.Tensor],
         backbone: MODEL_TYPE | None = None,
         optimizer: OptimizerCallable = optim.Adam,
@@ -36,6 +37,8 @@ class HeadModule(module.ModelModule):
 
         Args:
             head: The neural network that would be trained on the features.
+                If its a dictionary, it will be parsed to an object during the
+                `configure_model` step.
             criterion: The loss function to use.
             backbone: The feature extractor. If `None`, it will be expected
                 that the input batch returns the features directly.
@@ -48,15 +51,23 @@ class HeadModule(module.ModelModule):
         """
         super().__init__(metrics=metrics, postprocess=postprocess)
 
-        self.head = head
+        self.head = head  # type: ignore
         self.criterion = criterion
         self.backbone = backbone
         self.optimizer = optimizer
         self.lr_scheduler = lr_scheduler
 
+    @override
+    def configure_model(self) -> Any:
+        if self.backbone is not None:
+            grad.deactivate_requires_grad(self.backbone)
+
+        if isinstance(self.head, dict):
+            self.head: MODEL_TYPE = parser.parse_object(self.head, expected_type=nn.Module)
+
     @override
     def configure_optimizers(self) -> Any:
-        parameters = list(self.head.parameters())
+        parameters = self.head.parameters()
         optimizer = self.optimizer(parameters)
         lr_scheduler = self.lr_scheduler(optimizer)
         return {"optimizer": optimizer, "lr_scheduler": lr_scheduler}
@@ -66,11 +77,6 @@ class HeadModule(module.ModelModule):
         features = tensor if self.backbone is None else self.backbone(tensor)
         return self.head(features).squeeze(-1)
 
-    @override
-    def on_fit_start(self) -> None:
-        if self.backbone is not None:
-            grad.deactivate_requires_grad(self.backbone)
-
     @override
     def training_step(self, batch: INPUT_BATCH, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
         return self._batch_step(batch)
@@ -88,11 +94,6 @@ class HeadModule(module.ModelModule):
         tensor = INPUT_BATCH(*batch).data
         return tensor if self.backbone is None else self.backbone(tensor)
 
-    @override
-    def on_fit_end(self) -> None:
-        if self.backbone is not None:
-            grad.activate_requires_grad(self.backbone)
-
     def _batch_step(self, batch: INPUT_BATCH) -> STEP_OUTPUT:
         """Performs a model forward step and calculates the loss.