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

eva/core/loggers/log/parameters.py

@@ -0,0 +1,64 @@
+"""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,
+    )
+
+
+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/models/modules/head.py

@@ -54,9 +54,14 @@ class HeadModule(module.ModelModule):
         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)
+
     @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 +71,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 +88,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.
 

eva/core/models/modules/module.py

@@ -4,6 +4,7 @@ from typing import Any, Mapping
 
 import lightning.pytorch as pl
 import torch
+from lightning.pytorch.strategies.single_device import SingleDeviceStrategy
 from lightning.pytorch.utilities import memory
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from typing_extensions import override
@@ -46,6 +47,21 @@ class ModelModule(pl.LightningModule):
         """The default post-processes."""
         return batch_postprocess.BatchPostProcess()
 
+    @property
+    def metrics_device(self) -> torch.device:
+        """Returns the device by which the metrics should be calculated.
+
+        We allocate the metrics to CPU when operating on single device, as
+        it is much faster, but to GPU when employing multiple ones, as DDP
+        strategy requires the metrics to be allocated to the module's GPU.
+        """
+        move_to_cpu = isinstance(self.trainer.strategy, SingleDeviceStrategy)
+        return torch.device("cpu") if move_to_cpu else self.device
+
+    @override
+    def on_fit_start(self) -> None:
+        self.metrics.to(device=self.metrics_device)
+
     @override
     def on_train_batch_end(
         self,
@@ -59,6 +75,10 @@ class ModelModule(pl.LightningModule):
             batch_outputs=outputs,
         )
 
+    @override
+    def on_validation_start(self) -> None:
+        self.metrics.to(device=self.metrics_device)
+
     @override
     def on_validation_batch_end(
         self,
@@ -78,6 +98,10 @@ class ModelModule(pl.LightningModule):
     def on_validation_epoch_end(self) -> None:
         self._compute_and_log_metrics(self.metrics.validation_metrics)
 
+    @override
+    def on_test_start(self) -> None:
+        self.metrics.to(device=self.metrics_device)
+
     @override
     def on_test_batch_end(
         self,
@@ -110,7 +134,7 @@ class ModelModule(pl.LightningModule):
             The updated outputs.
         """
         self._postprocess(outputs)
-        return memory.recursive_detach(outputs, to_cpu=self.device.type == "cpu")
+        return memory.recursive_detach(outputs, to_cpu=self.metrics_device.type == "cpu")
 
     def _forward_and_log_metrics(
         self,

eva/core/trainers/_recorder.py

@@ -5,18 +5,41 @@ import json
 import os
 import statistics
 import sys
-from typing import Any, Dict, List, Mapping
+from typing import Dict, List, Mapping, TypedDict
 
 from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT
 from lightning_fabric.utilities import cloud_io
 from loguru import logger
 from omegaconf import OmegaConf
+from rich import console as rich_console
+from rich import table as rich_table
 from toolz import dicttoolz
 
 SESSION_METRICS = Mapping[str, List[float]]
 """Session metrics type-hint."""
 
 
+class SESSION_STATISTICS(TypedDict):
+    """Type-hint for aggregated metrics of multiple runs with mean & stdev."""
+
+    mean: float
+    stdev: float
+    values: List[float]
+
+
+class STAGE_RESULTS(TypedDict):
+    """Type-hint for metrics statstics for val & test stages."""
+
+    val: List[Dict[str, SESSION_STATISTICS]]
+    test: List[Dict[str, SESSION_STATISTICS]]
+
+
+class RESULTS_DICT(TypedDict):
+    """Type-hint for the final results dictionary."""
+
+    metrics: STAGE_RESULTS
+
+
 class SessionRecorder:
     """Multi-run (session) summary logger."""
 
@@ -25,6 +48,7 @@ class SessionRecorder:
         output_dir: str,
         results_file: str = "results.json",
         config_file: str = "config.yaml",
+        verbose: bool = True,
     ) -> None:
         """Initializes the recorder.
 
@@ -32,10 +56,12 @@ class SessionRecorder:
             output_dir: The destination folder to save the results.
             results_file: The name of the results json file.
             config_file: The name of the yaml configuration file.
+            verbose: Whether to print the session metrics.
         """
         self._output_dir = output_dir
         self._results_file = results_file
         self._config_file = config_file
+        self._verbose = verbose
 
         self._validation_metrics: List[SESSION_METRICS] = []
         self._test_metrics: List[SESSION_METRICS] = []
@@ -67,13 +93,13 @@ class SessionRecorder:
         self._update_validation_metrics(validation_scores)
         self._update_test_metrics(test_scores)
 
-    def compute(self) -> Dict[str, List[Dict[str, Any]]]:
+    def compute(self) -> STAGE_RESULTS:
         """Computes and returns the session statistics."""
         validation_statistics = list(map(_calculate_statistics, self._validation_metrics))
         test_statistics = list(map(_calculate_statistics, self._test_metrics))
         return {"val": validation_statistics, "test": test_statistics}
 
-    def export(self) -> Dict[str, Any]:
+    def export(self) -> RESULTS_DICT:
         """Exports the results."""
         statistics = self.compute()
         return {"metrics": statistics}
@@ -83,6 +109,8 @@ class SessionRecorder:
         results = self.export()
         _save_json(results, self.filename)
         self._save_config()
+        if self._verbose:
+            _print_results(results)
 
     def reset(self) -> None:
         """Resets the state of the tracked metrics."""
@@ -125,10 +153,10 @@ def _init_session_metrics(n_datasets: int) -> List[SESSION_METRICS]:
     return [collections.defaultdict(list) for _ in range(n_datasets)]
 
 
-def _calculate_statistics(session_metrics: SESSION_METRICS) -> Dict[str, float | List[float]]:
+def _calculate_statistics(session_metrics: SESSION_METRICS) -> Dict[str, SESSION_STATISTICS]:
     """Calculate the metric statistics of a dataset session run."""
 
-    def _calculate_metric_statistics(values: List[float]) -> Dict[str, float | List[float]]:
+    def _calculate_metric_statistics(values: List[float]) -> SESSION_STATISTICS:
         """Calculates and returns the metric statistics."""
         mean = statistics.mean(values)
         stdev = statistics.stdev(values) if len(values) > 1 else 0
@@ -137,7 +165,7 @@ def _calculate_statistics(session_metrics: SESSION_METRICS) -> Dict[str, float |
     return dicttoolz.valmap(_calculate_metric_statistics, session_metrics)
 
 
-def _save_json(data: Dict[str, Any], save_as: str = "data.json"):
+def _save_json(data: RESULTS_DICT, save_as: str = "data.json"):
     """Saves data to a json file."""
     if not save_as.endswith(".json"):
         raise ValueError()
@@ -146,4 +174,38 @@ def _save_json(data: Dict[str, Any], save_as: str = "data.json"):
     fs = cloud_io.get_filesystem(output_dir, anon=False)
     fs.makedirs(output_dir, exist_ok=True)
     with fs.open(save_as, "w") as file:
-        json.dump(data, file, indent=4, sort_keys=True)
+        json.dump(data, file, indent=2, sort_keys=True)
+
+
+def _print_results(results: RESULTS_DICT) -> None:
+    """Prints the results to the console."""
+    try:
+        for stage in ["val", "test"]:
+            for dataset_idx in range(len(results["metrics"][stage])):
+                _print_table(results["metrics"][stage][dataset_idx], stage, dataset_idx)
+    except Exception as e:
+        logger.error(f"Failed to print the results: {e}")
+
+
+def _print_table(metrics_dict: Dict[str, SESSION_STATISTICS], stage: str, dataset_idx: int):
+    """Prints the metrics of a single dataset as a table."""
+    metrics_table = rich_table.Table(
+        title=f"\n{stage.capitalize()} Dataset {dataset_idx}", title_style="bold"
+    )
+    metrics_table.add_column("Metric", style="cyan")
+    metrics_table.add_column("Mean", style="magenta")
+    metrics_table.add_column("Stdev", style="magenta")
+    metrics_table.add_column("All", style="magenta")
+
+    n_runs = len(metrics_dict[next(iter(metrics_dict))]["values"])
+    for metric_name, metric_dict in metrics_dict.items():
+        row = [
+            metric_name,
+            f'{metric_dict["mean"]:.3f}',
+            f'{metric_dict["stdev"]:.3f}',
+            ", ".join(f'{metric_dict["values"][i]:.3f}' for i in range(n_runs)),
+        ]
+        metrics_table.add_row(*row)
+
+    console = rich_console.Console()
+    console.print(metrics_table)

eva/core/trainers/functional.py

@@ -16,6 +16,7 @@ def run_evaluation_session(
     datamodule: datamodules.DataModule,
     *,
     n_runs: int = 1,
+    verbose: bool = True,
 ) -> None:
     """Runs a downstream evaluation session out-of-place.
 
@@ -29,11 +30,17 @@ def run_evaluation_session(
         base_model: The base model module to use.
         datamodule: The data module.
         n_runs: The amount of runs (fit and evaluate) to perform.
+        verbose: Whether to verbose the session metrics instead of
+            these of each individual runs and vice-versa.
     """
-    recorder = _recorder.SessionRecorder(output_dir=base_trainer.default_log_dir)
+    recorder = _recorder.SessionRecorder(output_dir=base_trainer.default_log_dir, verbose=verbose)
     for run_index in range(n_runs):
         validation_scores, test_scores = run_evaluation(
-            base_trainer, base_model, datamodule, run_id=f"run_{run_index}"
+            base_trainer,
+            base_model,
+            datamodule,
+            run_id=f"run_{run_index}",
+            verbose=not verbose,
         )
         recorder.update(validation_scores, test_scores)
     recorder.save()
@@ -45,6 +52,7 @@ def run_evaluation(
     datamodule: datamodules.DataModule,
     *,
     run_id: str | None = None,
+    verbose: bool = True,
 ) -> Tuple[_EVALUATE_OUTPUT, _EVALUATE_OUTPUT | None]:
     """Fits and evaluates a model out-of-place.
 
@@ -54,19 +62,22 @@ def run_evaluation(
         datamodule: The data module.
         run_id: The run id to be appended to the output log directory.
             If `None`, it will use the log directory of the trainer as is.
+        verbose: Whether to print the validation and test metrics
+            in the end of the training.
 
     Returns:
         A tuple of with the validation and the test metrics (if exists).
     """
     trainer, model = _utils.clone(base_trainer, base_model)
     trainer.setup_log_dirs(run_id or "")
-    return fit_and_validate(trainer, model, datamodule)
+    return fit_and_validate(trainer, model, datamodule, verbose=verbose)
 
 
 def fit_and_validate(
     trainer: eva_trainer.Trainer,
     model: modules.ModelModule,
     datamodule: datamodules.DataModule,
+    verbose: bool = True,
 ) -> Tuple[_EVALUATE_OUTPUT, _EVALUATE_OUTPUT | None]:
     """Fits and evaluates a model in-place.
 
@@ -77,13 +88,19 @@ def fit_and_validate(
         trainer: The trainer module to use and update in-place.
         model: The model module to use and update in-place.
         datamodule: The data module.
+        verbose: Whether to print the validation and test metrics
+            in the end of the training.
 
     Returns:
         A tuple of with the validation and the test metrics (if exists).
     """
     trainer.fit(model, datamodule=datamodule)
-    validation_scores = trainer.validate(datamodule=datamodule)
-    test_scores = None if datamodule.datasets.test is None else trainer.test(datamodule=datamodule)
+    validation_scores = trainer.validate(datamodule=datamodule, verbose=verbose)
+    test_scores = (
+        None
+        if datamodule.datasets.test is None
+        else trainer.test(datamodule=datamodule, verbose=verbose)
+    )
     return validation_scores, test_scores
 
 

eva/core/trainers/trainer.py

@@ -3,11 +3,14 @@
 import os
 from typing import Any
 
+import loguru
 from lightning.pytorch import loggers as pl_loggers
 from lightning.pytorch import trainer as pl_trainer
 from lightning.pytorch.utilities import argparse
+from lightning_fabric.utilities import cloud_io
 from typing_extensions import override
 
+from eva.core import loggers as eva_loggers
 from eva.core.data import datamodules
 from eva.core.models import modules
 from eva.core.trainers import _logging, functional
@@ -65,13 +68,23 @@ class Trainer(pl_trainer.Trainer):
             subdirectory: Whether to append a subdirectory to the output log.
         """
         self._log_dir = os.path.join(self.default_root_dir, self._session_id, subdirectory)
-        os.fspath(self._log_dir)
 
-        for logger in self.loggers:
-            if isinstance(logger, (pl_loggers.CSVLogger, pl_loggers.TensorBoardLogger)):
-                logger._root_dir = self.default_root_dir
-                logger._name = self._session_id
-                logger._version = subdirectory
+        enabled_loggers = []
+        if isinstance(self.loggers, list) and len(self.loggers) > 0:
+            for logger in self.loggers:
+                if isinstance(logger, (pl_loggers.CSVLogger, pl_loggers.TensorBoardLogger)):
+                    if not cloud_io._is_local_file_protocol(self.default_root_dir):
+                        loguru.logger.warning(
+                            f"Skipped {type(logger).__name__} as remote storage is not supported."
+                        )
+                        continue
+                    else:
+                        logger._root_dir = self.default_root_dir
+                        logger._name = self._session_id
+                        logger._version = subdirectory
+                enabled_loggers.append(logger)
+
+        self._loggers = enabled_loggers or [eva_loggers.DummyLogger(self._log_dir)]
 
     def run_evaluation_session(
         self,
@@ -94,4 +107,5 @@ class Trainer(pl_trainer.Trainer):
             base_model=model,
             datamodule=datamodule,
             n_runs=self._n_runs,
+            verbose=self._n_runs > 1,
         )

eva/vision/data/datasets/__init__.py

@@ -1,12 +1,6 @@
 """Vision Datasets API."""
 
-from eva.vision.data.datasets.classification import (
-    BACH,
-    CRC,
-    MHIST,
-    PatchCamelyon,
-    TotalSegmentatorClassification,
-)
+from eva.vision.data.datasets.classification import BACH, CRC, MHIST, PatchCamelyon
 from eva.vision.data.datasets.segmentation import ImageSegmentation, TotalSegmentator2D
 from eva.vision.data.datasets.vision import VisionDataset
 
@@ -16,7 +10,6 @@ __all__ = [
     "MHIST",
     "ImageSegmentation",
     "PatchCamelyon",
-    "TotalSegmentatorClassification",
     "TotalSegmentator2D",
     "VisionDataset",
 ]

eva/vision/data/datasets/_utils.py

@@ -1,6 +1,6 @@
 """Dataset related function and helper functions."""
 
-from typing import List, Tuple
+from typing import List, Sequence, Tuple
 
 
 def indices_to_ranges(indices: List[int]) -> List[Tuple[int, int]]:
@@ -33,11 +33,11 @@ def indices_to_ranges(indices: List[int]) -> List[Tuple[int, int]]:
     return ranges
 
 
-def ranges_to_indices(ranges: List[Tuple[int, int]]) -> List[int]:
+def ranges_to_indices(ranges: Sequence[Tuple[int, int]]) -> List[int]:
     """Unpacks a list of ranges to individual indices.
 
     Args:
-        ranges: The list of ranges to produce the indices from.
+        ranges: A sequence of ranges to produce the indices from.
 
     Return:
         A list of the indices.

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

@@ -4,12 +4,5 @@ from eva.vision.data.datasets.classification.bach import BACH
 from eva.vision.data.datasets.classification.crc import CRC
 from eva.vision.data.datasets.classification.mhist import MHIST
 from eva.vision.data.datasets.classification.patch_camelyon import PatchCamelyon
-from eva.vision.data.datasets.classification.total_segmentator import TotalSegmentatorClassification
 
-__all__ = [
-    "BACH",
-    "CRC",
-    "MHIST",
-    "PatchCamelyon",
-    "TotalSegmentatorClassification",
-]
+__all__ = ["BACH", "CRC", "MHIST", "PatchCamelyon"]

eva/vision/data/datasets/segmentation/base.py

@@ -3,38 +3,28 @@
 import abc
 from typing import Any, Callable, Dict, List, Tuple
 
-import numpy as np
+from torchvision import tv_tensors
 from typing_extensions import override
 
 from eva.vision.data.datasets import vision
 
 
-class ImageSegmentation(vision.VisionDataset[Tuple[np.ndarray, np.ndarray]], abc.ABC):
+class ImageSegmentation(vision.VisionDataset[Tuple[tv_tensors.Image, tv_tensors.Mask]], abc.ABC):
     """Image segmentation abstract dataset."""
 
     def __init__(
         self,
-        image_transforms: Callable | None = None,
-        target_transforms: Callable | None = None,
-        image_target_transforms: Callable | None = None,
+        transforms: Callable | None = None,
     ) -> None:
         """Initializes the image segmentation base class.
 
         Args:
-            image_transforms: A function/transform that takes in an image
-                and returns a transformed version.
-            target_transforms: A function/transform that takes in the target
-                and transforms it.
-            image_target_transforms: A function/transforms that takes in an
+            transforms: A function/transforms that takes in an
                 image and a label and returns the transformed versions of both.
-                This transform happens after the `image_transforms` and
-                `target_transforms`.
         """
         super().__init__()
 
-        self._image_transforms = image_transforms
-        self._target_transforms = target_transforms
-        self._image_target_transforms = image_target_transforms
+        self._transforms = transforms
 
     @property
     def classes(self) -> List[str] | None:
@@ -56,25 +46,26 @@ class ImageSegmentation(vision.VisionDataset[Tuple[np.ndarray, np.ndarray]], abc
         """
 
     @abc.abstractmethod
-    def load_image(self, index: int) -> np.ndarray:
+    def load_image(self, index: int) -> tv_tensors.Image:
         """Loads and returns the `index`'th image sample.
 
         Args:
             index: The index of the data sample to load.
 
         Returns:
-            The image as a numpy array.
+            An image torchvision tensor (channels, height, width).
         """
 
     @abc.abstractmethod
-    def load_mask(self, index: int) -> np.ndarray:
-        """Returns the `index`'th target mask sample.
+    def load_mask(self, index: int) -> tv_tensors.Mask:
+        """Returns the `index`'th target masks sample.
 
         Args:
-            index: The index of the data sample target mask to load.
+            index: The index of the data sample target masks to load.
 
         Returns:
-            The sample mask as a stack of binary mask arrays (label, height, width).
+            The semantic mask as a (H x W) shaped tensor with integer
+            values which represent the pixel class id.
         """
 
     @abc.abstractmethod
@@ -83,30 +74,24 @@ class ImageSegmentation(vision.VisionDataset[Tuple[np.ndarray, np.ndarray]], abc
         raise NotImplementedError
 
     @override
-    def __getitem__(self, index: int) -> Tuple[np.ndarray, np.ndarray]:
+    def __getitem__(self, index: int) -> Tuple[tv_tensors.Image, tv_tensors.Mask]:
         image = self.load_image(index)
         mask = self.load_mask(index)
         return self._apply_transforms(image, mask)
 
     def _apply_transforms(
-        self, image: np.ndarray, target: np.ndarray
-    ) -> Tuple[np.ndarray, np.ndarray]:
+        self, image: tv_tensors.Image, mask: tv_tensors.Mask
+    ) -> Tuple[tv_tensors.Image, tv_tensors.Mask]:
         """Applies the transforms to the provided data and returns them.
 
         Args:
             image: The desired image.
-            target: The target of the image.
+            mask: The target segmentation mask.
 
         Returns:
-            A tuple with the image and the target transformed.
+            A tuple with the image and the masks transformed.
         """
-        if self._image_transforms is not None:
-            image = self._image_transforms(image)
+        if self._transforms is not None:
+            image, mask = self._transforms(image, mask)
 
-        if self._target_transforms is not None:
-            target = self._target_transforms(target)
-
-        if self._image_target_transforms is not None:
-            image, target = self._image_target_transforms(image, target)
-
-        return image, target
+        return image, mask