kaiko-eva 0.1.8__py3-none-any.whl → 0.2.1__py3-none-any.whl

eva/core/data/datasets/base.py

@@ -1,6 +1,7 @@
 """Base dataset class."""
 
 import abc
+from typing import Generic, TypeVar
 
 from eva.core.data.datasets import dataset
 
@@ -55,11 +56,15 @@ class Dataset(dataset.TorchDataset):
         """
 
 
-class MapDataset(Dataset):
+DataSample = TypeVar("DataSample")
+"""The data sample type."""
+
+
+class MapDataset(Dataset, abc.ABC, Generic[DataSample]):
     """Abstract base class for all map-style datasets."""
 
     @abc.abstractmethod
-    def __getitem__(self, index: int):
+    def __getitem__(self, index: int) -> DataSample:
         """Retrieves the item at the given index.
 
         Args:

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

@@ -12,7 +12,7 @@ class EmbeddingsClassificationDataset(embeddings_base.EmbeddingsDataset[torch.Te
     """Embeddings dataset class for classification tasks."""
 
     @override
-    def _load_embeddings(self, index: int) -> torch.Tensor:
+    def load_embeddings(self, index: int) -> torch.Tensor:
         filename = self.filename(index)
         embeddings_path = os.path.join(self._root, filename)
         tensor = torch.load(embeddings_path, map_location="cpu")
@@ -25,7 +25,7 @@ class EmbeddingsClassificationDataset(embeddings_base.EmbeddingsDataset[torch.Te
         return tensor.squeeze(0)
 
     @override
-    def _load_target(self, index: int) -> torch.Tensor:
+    def load_target(self, index: int) -> torch.Tensor:
         target = self._data.at[index, self._column_mapping["target"]]
         return torch.tensor(target, dtype=torch.int64)
 

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

@@ -66,7 +66,7 @@ class MultiEmbeddingsClassificationDataset(embeddings_base.EmbeddingsDataset[tor
         self._multi_ids = list(self._data[self._column_mapping["multi_id"]].unique())
 
     @override
-    def _load_embeddings(self, index: int) -> torch.Tensor:
+    def load_embeddings(self, index: int) -> torch.Tensor:
         """Loads and stacks all embedding corresponding to the `index`'th multi_id."""
         # Get all embeddings for the given index (multi_id)
         multi_id = self._multi_ids[index]
@@ -89,7 +89,7 @@ class MultiEmbeddingsClassificationDataset(embeddings_base.EmbeddingsDataset[tor
         return embeddings
 
     @override
-    def _load_target(self, index: int) -> np.ndarray:
+    def load_target(self, index: int) -> np.ndarray:
         """Returns the target corresponding to the `index`'th multi_id.
 
         This method assumes that all the embeddings corresponding to the same `multi_id`

eva/core/data/datasets/embeddings.py

@@ -98,12 +98,12 @@ class EmbeddingsDataset(base.Dataset, Generic[TargetType]):
         Returns:
             A data sample and its target.
         """
-        embeddings = self._load_embeddings(index)
-        target = self._load_target(index)
+        embeddings = self.load_embeddings(index)
+        target = self.load_target(index)
         return self._apply_transforms(embeddings, target)
 
     @abc.abstractmethod
-    def _load_embeddings(self, index: int) -> torch.Tensor:
+    def load_embeddings(self, index: int) -> torch.Tensor:
         """Returns the `index`'th embedding sample.
 
         Args:
@@ -114,7 +114,7 @@ class EmbeddingsDataset(base.Dataset, Generic[TargetType]):
         """
 
     @abc.abstractmethod
-    def _load_target(self, index: int) -> TargetType:
+    def load_target(self, index: int) -> TargetType:
         """Returns the `index`'th target sample.
 
         Args:

eva/core/data/samplers/classification/balanced.py

@@ -4,6 +4,7 @@ from collections import defaultdict
 from typing import Dict, Iterator, List
 
 import numpy as np
+from loguru import logger
 from typing_extensions import override
 
 from eva.core.data import datasets
@@ -33,6 +34,7 @@ class BalancedSampler(SamplerWithDataSource[int]):
         self._replacement = replacement
         self._class_indices: Dict[int, List[int]] = defaultdict(list)
         self._random_generator = np.random.default_rng(seed)
+        self._indices: List[int] = []
 
     def __len__(self) -> int:
         """Returns the total number of samples."""
@@ -44,18 +46,7 @@ class BalancedSampler(SamplerWithDataSource[int]):
         Returns:
             Iterator yielding dataset indices.
         """
-        indices = []
-
-        for class_idx in self._class_indices:
-            class_indices = self._class_indices[class_idx]
-            sampled_indices = self._random_generator.choice(
-                class_indices, size=self._num_samples, replace=self._replacement
-            ).tolist()
-            indices.extend(sampled_indices)
-
-        self._random_generator.shuffle(indices)
-
-        return iter(indices)
+        return iter(self._indices)
 
     @override
     def set_dataset(self, data_source: datasets.MapDataset):
@@ -72,13 +63,13 @@ class BalancedSampler(SamplerWithDataSource[int]):
         self._make_indices()
 
     def _make_indices(self):
-        """Builds indices for each class in the dataset."""
+        """Samples the indices for each class in the dataset."""
         self._class_indices.clear()
-
-        for idx in tqdm(
-            range(len(self.data_source)), desc="Fetching class indices for balanced sampler"
-        ):
-            _, target, _ = DataSample(*self.data_source[idx])
+        for idx in tqdm(range(len(self.data_source)), desc="Fetching class indices for sampler"):
+            if hasattr(self.data_source, "load_target"):
+                target = self.data_source.load_target(idx)  # type: ignore
+            else:
+                _, target, _ = DataSample(*self.data_source[idx])
             if target is None:
                 raise ValueError("The dataset must return non-empty targets.")
             if target.numel() != 1:
@@ -94,3 +85,13 @@ class BalancedSampler(SamplerWithDataSource[int]):
                         f"Class {class_idx} has only {len(indices)} samples, "
                         f"which is less than the required {self._num_samples} samples."
                     )
+
+        self._indices = []
+        for class_idx in self._class_indices:
+            class_indices = self._class_indices[class_idx]
+            sampled_indices = self._random_generator.choice(
+                class_indices, size=self._num_samples, replace=self._replacement
+            ).tolist()
+            self._indices.extend(sampled_indices)
+        self._random_generator.shuffle(self._indices)
+        logger.debug(f"Sampled indices: {self._indices}")

eva/core/loggers/utils/wandb.py

@@ -0,0 +1,33 @@
+# type: ignore
+"""Utility functions for logging with Weights & Biases."""
+
+from typing import Any, Dict
+
+from loguru import logger
+
+
+def rename_active_run(name: str) -> None:
+    """Renames the current run."""
+    import wandb
+
+    if wandb.run:
+        wandb.run.name = name
+        wandb.run.save()
+    else:
+        logger.warning("No active wandb run found that could be renamed.")
+
+
+def init_run(name: str, init_kwargs: Dict[str, Any]) -> None:
+    """Initializes a new run. If there is an active run, it will be renamed and reused."""
+    import wandb
+
+    init_kwargs["name"] = name
+    rename_active_run(name)
+    wandb.init(**init_kwargs)
+
+
+def finish_run() -> None:
+    """Finish the current run."""
+    import wandb
+
+    wandb.finish()

eva/core/models/modules/head.py

@@ -1,6 +1,6 @@
-""""Neural Network Head Module."""
+"""Neural Network Head Module."""
 
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, List
 
 import torch
 from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
@@ -108,7 +108,9 @@ class HeadModule(module.ModelModule):
         return self._batch_step(batch)
 
     @override
-    def predict_step(self, batch: INPUT_BATCH, *args: Any, **kwargs: Any) -> torch.Tensor:
+    def predict_step(
+        self, batch: INPUT_BATCH, *args: Any, **kwargs: Any
+    ) -> torch.Tensor | List[torch.Tensor]:
         tensor = INPUT_BATCH(*batch).data
         return tensor if self.backbone is None else self.backbone(tensor)
 

eva/core/models/modules/typings.py

@@ -1,6 +1,6 @@
 """Type annotations for model modules."""
 
-from typing import Any, Dict, NamedTuple
+from typing import Any, Dict, List, NamedTuple
 
 import lightning.pytorch as pl
 import torch
@@ -13,7 +13,7 @@ MODEL_TYPE = nn.Module | pl.LightningModule
 class INPUT_BATCH(NamedTuple):
     """The default input batch data scheme."""
 
-    data: torch.Tensor
+    data: torch.Tensor | List[torch.Tensor]
     """The data batch."""
 
     targets: torch.Tensor | None = None

eva/core/models/transforms/__init__.py

@@ -1,6 +1,7 @@
 """Model outputs transforms API."""
 
+from eva.core.models.transforms.as_discrete import AsDiscrete
 from eva.core.models.transforms.extract_cls_features import ExtractCLSFeatures
 from eva.core.models.transforms.extract_patch_features import ExtractPatchFeatures
 
-__all__ = ["ExtractCLSFeatures", "ExtractPatchFeatures"]
+__all__ = ["AsDiscrete", "ExtractCLSFeatures", "ExtractPatchFeatures"]

eva/core/models/transforms/as_discrete.py

@@ -0,0 +1,57 @@
+"""Defines the AsDiscrete transformation."""
+
+import torch
+
+
+class AsDiscrete:
+    """Convert the logits tensor to discrete values."""
+
+    def __init__(
+        self,
+        argmax: bool = False,
+        to_onehot: int | bool | None = None,
+        threshold: float | None = None,
+    ) -> None:
+        """Convert the input tensor/array into discrete values.
+
+        Args:
+            argmax: Whether to execute argmax function on input data before transform.
+            to_onehot: if not None, convert input data into the one-hot format with
+                specified number of classes. If bool, it will try to infer the number
+                of classes.
+            threshold: If not None, threshold the float values to int number 0 or 1
+                with specified threshold.
+        """
+        super().__init__()
+
+        self._argmax = argmax
+        self._to_onehot = to_onehot
+        self._threshold = threshold
+
+    def __call__(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Call method for the transformation."""
+        if self._argmax:
+            tensor = torch.argmax(tensor, dim=1, keepdim=True)
+
+        if self._to_onehot is not None:
+            tensor = _one_hot(tensor, num_classes=self._to_onehot, dim=1, dtype=torch.long)
+
+        if self._threshold is not None:
+            tensor = tensor >= self._threshold
+
+        return tensor
+
+
+def _one_hot(
+    tensor: torch.Tensor, num_classes: int, dtype: torch.dtype = torch.float, dim: int = 1
+) -> torch.Tensor:
+    """Convert input tensor into one-hot format (implementation taken from MONAI)."""
+    shape = list(tensor.shape)
+    if shape[dim] != 1:
+        raise AssertionError(f"Input tensor must have 1 channel at dim {dim}.")
+
+    shape[dim] = num_classes
+    o = torch.zeros(size=shape, dtype=dtype, device=tensor.device)
+    tensor = o.scatter_(dim=dim, index=tensor.long(), value=1)
+
+    return tensor

eva/core/models/wrappers/_utils.py

@@ -1,8 +1,17 @@
 """Utilities and helper functions for models."""
 
+import hashlib
+import os
+import sys
+from typing import Any, Dict
+
+import torch
+from fsspec.core import url_to_fs
 from lightning_fabric.utilities import cloud_io
 from loguru import logger
-from torch import nn
+from torch import hub, nn
+
+from eva.core.utils.progress_bar import tqdm
 
 
 def load_model_weights(model: nn.Module, checkpoint_path: str) -> None:
@@ -23,3 +32,114 @@ def load_model_weights(model: nn.Module, checkpoint_path: str) -> None:
         model.load_state_dict(checkpoint, strict=True)
 
     logger.info(f"Loading weights from '{checkpoint_path}' completed successfully.")
+
+
+def load_state_dict_from_url(
+    url: str,
+    *,
+    model_dir: str | None = None,
+    filename: str | None = None,
+    progress: bool = True,
+    md5: str | None = None,
+    force: bool = False,
+) -> Dict[str, Any]:
+    """Loads the Torch serialized object at the given URL.
+
+    If the object is already present and valid in `model_dir`, it's
+    deserialized and returned.
+
+    The default value of ``model_dir`` is ``<hub_dir>/checkpoints`` where
+    ``hub_dir`` is the directory returned by :func:`~torch.hub.get_dir`.
+
+    Args:
+        url: URL of the object to download.
+        model_dir: Directory in which to save the object.
+        filename: Name for the downloaded file. Filename from ``url`` will be used if not set.
+        progress: Whether or not to display a progress bar to stderr.
+        md5: MD5 file code to check whether the file is valid. If not, it will re-download it.
+        force: Whether to download the file regardless if it exists.
+    """
+    model_dir = model_dir or os.path.join(hub.get_dir(), "checkpoints")
+    os.makedirs(model_dir, exist_ok=True)
+
+    cached_file = os.path.join(model_dir, filename or os.path.basename(url))
+    if force or not os.path.exists(cached_file) or not _check_integrity(cached_file, md5):
+        sys.stderr.write(f"Downloading: '{url}' to {cached_file}\n")
+        _download_url_to_file(url, cached_file, progress=progress)
+        if md5 is None or not _check_integrity(cached_file, md5):
+            sys.stderr.write(f"File MD5: {_calculate_md5(cached_file)}\n")
+
+    return torch.load(cached_file, map_location="cpu")
+
+
+def _download_url_to_file(
+    url: str,
+    dst: str,
+    *,
+    progress: bool = True,
+) -> None:
+    """Download object at the given URL to a local path.
+
+    Args:
+        url: URL of the object to download.
+        dst: Full path where object will be saved.
+        chunk_size: The size of each chunk to read in bytes.
+        progress: Whether or not to display a progress bar to stderr.
+    """
+    try:
+        _download_with_fsspec(url=url, dst=dst, progress=progress)
+    except Exception:
+        try:
+            hub.download_url_to_file(url=url, dst=dst, progress=progress)
+        except Exception as hub_e:
+            raise RuntimeError(
+                f"Failed to download file from {url} using both fsspec and hub."
+            ) from hub_e
+
+
+def _download_with_fsspec(
+    url: str,
+    dst: str,
+    *,
+    chunk_size: int = 1024 * 1024,
+    progress: bool = True,
+) -> None:
+    """Download object at the given URL to a local path using fsspec.
+
+    Args:
+        url: URL of the object to download.
+        dst: Full path where object will be saved.
+        chunk_size: The size of each chunk to read in bytes.
+        progress: Whether or not to display a progress bar to stderr.
+    """
+    filesystem, _ = url_to_fs(url, anon=False)
+    total_size_bytes = filesystem.size(url)
+    with (
+        filesystem.open(url, "rb") as remote_file,
+        tqdm(
+            total=total_size_bytes,
+            unit="iB",
+            unit_scale=True,
+            unit_divisor=1024,
+            disable=not progress,
+        ) as pbar,
+    ):
+        with open(dst, "wb") as local_file:
+            while True:
+                data = remote_file.read(chunk_size)
+                if not data:
+                    break
+
+                local_file.write(data)
+                pbar.update(chunk_size)
+
+
+def _calculate_md5(path: str) -> str:
+    """Calculate the md5 hash of a file."""
+    with open(path, "rb") as file:
+        return hashlib.md5(file.read(), usedforsecurity=False).hexdigest()
+
+
+def _check_integrity(path: str, md5: str | None) -> bool:
+    """Check if the file matches the specified md5 hash."""
+    return (md5 is None) or (md5 == _calculate_md5(path))

eva/core/trainers/functional.py

@@ -39,7 +39,7 @@ def run_evaluation_session(
             base_trainer,
             base_model,
             datamodule,
-            run_id=f"run_{run_index}",
+            run_id=run_index,
             verbose=not verbose,
         )
         recorder.update(validation_scores, test_scores)
@@ -51,7 +51,7 @@ def run_evaluation(
     base_model: modules.ModelModule,
     datamodule: datamodules.DataModule,
     *,
-    run_id: str | None = None,
+    run_id: int | None = None,
     verbose: bool = True,
 ) -> Tuple[_EVALUATE_OUTPUT, _EVALUATE_OUTPUT | None]:
     """Fits and evaluates a model out-of-place.
@@ -61,7 +61,6 @@ def run_evaluation(
         base_model: The model module to use but not modify.
         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.
 
@@ -70,8 +69,12 @@ def run_evaluation(
     """
     trainer, model = _utils.clone(base_trainer, base_model)
     model.configure_model()
-    trainer.setup_log_dirs(run_id or "")
-    return fit_and_validate(trainer, model, datamodule, verbose=verbose)
+
+    trainer.init_logger_run(run_id)
+    results = fit_and_validate(trainer, model, datamodule, verbose=verbose)
+    trainer.finish_logger_run(run_id)
+
+    return results
 
 
 def fit_and_validate(

eva/core/trainers/trainer.py

@@ -12,6 +12,7 @@ from typing_extensions import override
 
 from eva.core import loggers as eva_loggers
 from eva.core.data import datamodules
+from eva.core.loggers.utils import wandb as wandb_utils
 from eva.core.models import modules
 from eva.core.trainers import _logging, functional
 
@@ -53,7 +54,7 @@ class Trainer(pl_trainer.Trainer):
         self._session_id: str = _logging.generate_session_id()
         self._log_dir: str = self.default_log_dir
 
-        self.setup_log_dirs()
+        self.init_logger_run(0)
 
     @property
     def default_log_dir(self) -> str:
@@ -65,31 +66,45 @@ class Trainer(pl_trainer.Trainer):
     def log_dir(self) -> str | None:
         return self.strategy.broadcast(self._log_dir)
 
-    def setup_log_dirs(self, subdirectory: str = "") -> None:
-        """Setups the logging directory of the trainer and experimental loggers in-place.
+    def init_logger_run(self, run_id: int | None) -> None:
+        """Setup the loggers & log directories when starting a new run.
 
         Args:
-            subdirectory: Whether to append a subdirectory to the output log.
+            run_id: The id of the current run.
         """
+        subdirectory = f"run_{run_id}" if run_id is not None else ""
         self._log_dir = os.path.join(self.default_root_dir, self._session_id, 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)
+        for logger in self.loggers or []:
+            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
+            elif isinstance(logger, pl_loggers.WandbLogger):
+                task_name = self.default_root_dir.split("/")[-1]
+                run_name = os.getenv("WANDB_RUN_NAME", f"{task_name}_{self._session_id}")
+                wandb_utils.init_run(f"{run_name}_{run_id}", logger._wandb_init)
+            enabled_loggers.append(logger)
 
         self._loggers = enabled_loggers or [eva_loggers.DummyLogger(self._log_dir)]
 
+    def finish_logger_run(self, run_id: int | None) -> None:
+        """Finish the current run in the enabled loggers.
+
+        Args:
+            run_id: The id of the current run.
+        """
+        for logger in self.loggers or []:
+            if isinstance(logger, pl_loggers.WandbLogger):
+                wandb_utils.finish_run()
+
     def run_evaluation_session(
         self,
         model: modules.ModelModule,

eva/core/utils/suppress_logs.py

@@ -0,0 +1,28 @@
+"""Context manager to temporarily suppress all logging outputs."""
+
+import logging
+import sys
+from types import TracebackType
+from typing import Type
+
+
+class SuppressLogs:
+    """Context manager to suppress all logs but print exceptions if they occur."""
+
+    def __enter__(self) -> None:
+        """Temporarily increase log level to suppress all logs."""
+        self._logger = logging.getLogger()
+        self._previous_level = self._logger.level
+        self._logger.setLevel(logging.CRITICAL + 1)
+
+    def __exit__(
+        self,
+        exc_type: Type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool:
+        """Restores the previous logging level and print exceptions."""
+        self._logger.setLevel(self._previous_level)
+        if exc_value:
+            print(f"Error: {exc_value}", file=sys.stderr)
+        return False

eva/vision/data/__init__.py

@@ -1,5 +1,5 @@
 """Vision data API."""
 
-from eva.vision.data import datasets, transforms
+from eva.vision.data import datasets, transforms, tv_tensors
 
-__all__ = ["datasets", "transforms"]
+__all__ = ["datasets", "transforms", "tv_tensors"]

eva/vision/data/dataloaders/__init__.py

@@ -0,0 +1,5 @@
+"""Dataloader related utilities and functions."""
+
+from eva.vision.data.dataloaders import collate_fn
+
+__all__ = ["collate_fn"]

eva/vision/data/dataloaders/collate_fn/__init__.py

@@ -0,0 +1,5 @@
+"""Dataloader collate API."""
+
+from eva.vision.data.dataloaders.collate_fn.collection import collection_collate
+
+__all__ = ["collection_collate"]

eva/vision/data/dataloaders/collate_fn/collection.py

@@ -0,0 +1,22 @@
+"""Data only collate filter function."""
+
+from typing import Any, List
+
+import torch
+
+from eva.core.models.modules.typings import INPUT_BATCH
+
+
+def collection_collate(batch: List[List[INPUT_BATCH]]) -> Any:
+    """Collate function for stacking a collection of data samples.
+
+    Args:
+        batch: The batch to be collated.
+
+    Returns:
+        The collated batch.
+    """
+    tensors, targets, metadata = zip(*batch, strict=False)
+    batch_tensors = torch.cat(list(map(torch.stack, tensors)))
+    batch_targets = torch.cat(list(map(torch.stack, targets)))
+    return batch_tensors, batch_targets, metadata

eva/vision/data/datasets/__init__.py

@@ -2,19 +2,23 @@
 
 from eva.vision.data.datasets.classification import (
     BACH,
+    BRACS,
     CRC,
     MHIST,
     PANDA,
+    BreaKHis,
     Camelyon16,
+    GleasonArvaniti,
     PANDASmall,
     PatchCamelyon,
+    UniToPatho,
     WsiClassificationDataset,
 )
 from eva.vision.data.datasets.segmentation import (
     BCSS,
+    BTCV,
     CoNSeP,
     EmbeddingsSegmentationDataset,
-    ImageSegmentation,
     LiTS,
     LiTSBalanced,
     MoNuSAC,
@@ -25,17 +29,21 @@ from eva.vision.data.datasets.wsi import MultiWsiDataset, WsiDataset
 
 __all__ = [
     "BACH",
+    "BTCV",
     "BCSS",
+    "BreaKHis",
+    "BRACS",
     "CRC",
+    "GleasonArvaniti",
     "MHIST",
     "PANDA",
     "PANDASmall",
     "Camelyon16",
     "PatchCamelyon",
+    "UniToPatho",
     "WsiClassificationDataset",
     "CoNSeP",
     "EmbeddingsSegmentationDataset",
-    "ImageSegmentation",
     "LiTS",
     "LiTSBalanced",
     "MoNuSAC",