kaiko-eva 0.2.0__tar.gz → 0.2.2__tar.gz

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kaiko-eva
-Version: 0.2.0
+Version: 0.2.2
 Summary: Evaluation Framework for oncology foundation models.
 Keywords: machine-learning,evaluation-framework,oncology,foundation-models
 Author-Email: Ioannis Gatopoulos <ioannis@kaiko.ai>, =?utf-8?q?Nicolas_K=C3=A4nzig?= <nicolas@kaiko.ai>, Roman Moser <roman@kaiko.ai>
@@ -241,6 +241,7 @@ Requires-Dist: scikit-image>=0.24.0; extra == "vision"
 Requires-Dist: imagesize>=1.4.1; extra == "vision"
 Requires-Dist: scipy>=1.14.0; extra == "vision"
 Requires-Dist: monai>=1.3.2; extra == "vision"
+Requires-Dist: einops>=0.8.1; extra == "vision"
 Provides-Extra: all
 Requires-Dist: h5py>=3.10.0; extra == "all"
 Requires-Dist: nibabel>=4.0.1; extra == "all"
@@ -253,6 +254,7 @@ Requires-Dist: scikit-image>=0.24.0; extra == "all"
 Requires-Dist: imagesize>=1.4.1; extra == "all"
 Requires-Dist: scipy>=1.14.0; extra == "all"
 Requires-Dist: monai>=1.3.2; extra == "all"
+Requires-Dist: einops>=0.8.1; extra == "all"
 Description-Content-Type: text/markdown
 
 <div align="center">
@@ -291,7 +293,7 @@ Check out the [documentation](https://kaiko-ai.github.io/eva/) for more informat
 
 ### Highlights:
 - Easy and reliable benchmark of Oncology FMs
-- Supports path-level classification, slide-level classification and semantic segmentation downstream tasks
+- Supports patch-level classification, slide-level classification and semantic segmentation downstream tasks
 - Automatic embedding inference and evaluation of a downstream task
 - Native support of popular medical [datasets](https://kaiko-ai.github.io/eva/dev/datasets/) and models
 - Produce statistics over multiple evaluation fits and multiple metrics
@@ -446,26 +448,28 @@ input, resulting in a faster evaluation.
 Here are some examples to get you started:
 
 - Perform a downstream offline **classification** evaluation of `DINO ViT-S/16`
-on the `BACH` dataset with linear probing by first inferring the embeddings
-and then performing 5 sequential fits:
+on the `BACH` dataset with linear probing by first pre-calculating the embeddings:
   ```sh
-  export DOWNLOAD_DATA=true
-  eva predict_fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/dino_vit/offline/bach.yaml
+  DOWNLOAD_DATA=true \
+  MODEL_NAME=universal/vit_small_patch16_224_dino \
+  eva predict_fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/pathology/offline/classification/bach.yaml
   ```
 
-- Perform a downstream online **segmentation** evaluation of `DINO ViT-S/16` on the
-`MoNuSAC` dataset with the `ConvDecoderMS` decoder:
+- Perform a downstream online **segmentation** evaluation of `DINO ViT-S/16` on the `MoNuSAC` dataset with the `ConvDecoderWithImage` decoder:
   ```sh
-  export DOWNLOAD_DATA=true
-  eva fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/dino_vit/online/monusac.yaml
+  DOWNLOAD_DATA=true \
+  MODEL_NAME=universal/vit_small_patch16_224_dino \
+  eva fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/pathology/online/segmentation/monusac.yaml
   ```
 
+By default `eva` will perform 5 evaluation runs using different seeds, however, you can control the number of runs through the `N_RUNS` environment variable or in the configuration file. The results will be saved to `./logs` by default, or to `OUTPUT_ROOT` if specified.
+
 For more examples, take a look at the [configs](https://github.com/kaiko-ai/eva/tree/main/configs)
-and [tutorials](https://kaiko-ai.github.io/eva/dev/user-guide/advanced/replicate_evaluations/).
+and [tutorials](https://kaiko-ai.github.io/eva/main/user-guide/advanced/replicate_evaluations/).
 
 > [!NOTE]
 > All the datasets that support automatic download in the repo have by default the option to automatically download set to false.
-> For automatic download you have to manually set the environmental variable `DOWNLOAD_DATA=true` or in the configuration file `download=true`.
+> For automatic download you have to manually set the environment variable `DOWNLOAD_DATA=true` or in the configuration file `download=true`.
 
 ## Leaderboards
 

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/README.md

@@ -34,7 +34,7 @@ Check out the [documentation](https://kaiko-ai.github.io/eva/) for more informat
 
 ### Highlights:
 - Easy and reliable benchmark of Oncology FMs
-- Supports path-level classification, slide-level classification and semantic segmentation downstream tasks
+- Supports patch-level classification, slide-level classification and semantic segmentation downstream tasks
 - Automatic embedding inference and evaluation of a downstream task
 - Native support of popular medical [datasets](https://kaiko-ai.github.io/eva/dev/datasets/) and models
 - Produce statistics over multiple evaluation fits and multiple metrics
@@ -189,26 +189,28 @@ input, resulting in a faster evaluation.
 Here are some examples to get you started:
 
 - Perform a downstream offline **classification** evaluation of `DINO ViT-S/16`
-on the `BACH` dataset with linear probing by first inferring the embeddings
-and then performing 5 sequential fits:
+on the `BACH` dataset with linear probing by first pre-calculating the embeddings:
   ```sh
-  export DOWNLOAD_DATA=true
-  eva predict_fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/dino_vit/offline/bach.yaml
+  DOWNLOAD_DATA=true \
+  MODEL_NAME=universal/vit_small_patch16_224_dino \
+  eva predict_fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/pathology/offline/classification/bach.yaml
   ```
 
-- Perform a downstream online **segmentation** evaluation of `DINO ViT-S/16` on the
-`MoNuSAC` dataset with the `ConvDecoderMS` decoder:
+- Perform a downstream online **segmentation** evaluation of `DINO ViT-S/16` on the `MoNuSAC` dataset with the `ConvDecoderWithImage` decoder:
   ```sh
-  export DOWNLOAD_DATA=true
-  eva fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/dino_vit/online/monusac.yaml
+  DOWNLOAD_DATA=true \
+  MODEL_NAME=universal/vit_small_patch16_224_dino \
+  eva fit --config https://raw.githubusercontent.com/kaiko-ai/eva/main/configs/vision/pathology/online/segmentation/monusac.yaml
   ```
 
+By default `eva` will perform 5 evaluation runs using different seeds, however, you can control the number of runs through the `N_RUNS` environment variable or in the configuration file. The results will be saved to `./logs` by default, or to `OUTPUT_ROOT` if specified.
+
 For more examples, take a look at the [configs](https://github.com/kaiko-ai/eva/tree/main/configs)
-and [tutorials](https://kaiko-ai.github.io/eva/dev/user-guide/advanced/replicate_evaluations/).
+and [tutorials](https://kaiko-ai.github.io/eva/main/user-guide/advanced/replicate_evaluations/).
 
 > [!NOTE]
 > All the datasets that support automatic download in the repo have by default the option to automatically download set to false.
-> For automatic download you have to manually set the environmental variable `DOWNLOAD_DATA=true` or in the configuration file `download=true`.
+> For automatic download you have to manually set the environment variable `DOWNLOAD_DATA=true` or in the configuration file `download=true`.
 
 ## Leaderboards
 

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "kaiko-eva"
-version = "0.2.0"
+version = "0.2.2"
 description = "Evaluation Framework for oncology foundation models."
 keywords = [
     "machine-learning",
@@ -70,6 +70,7 @@ vision = [
     "imagesize>=1.4.1",
     "scipy>=1.14.0",
     "monai>=1.3.2",
+    "einops>=0.8.1",
 ]
 all = [
     "h5py>=3.10.0",
@@ -83,6 +84,7 @@ all = [
     "imagesize>=1.4.1",
     "scipy>=1.14.0",
     "monai>=1.3.2",
+    "einops>=0.8.1",
 ]
 
 [project.scripts]

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/src/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:

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/src/eva/core/models/modules/head.py

@@ -1,6 +1,6 @@
 """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)
 

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/src/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

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/src/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"]

kaiko_eva-0.2.2/src/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

kaiko_eva-0.2.2/src/eva/core/models/wrappers/_utils.py

@@ -0,0 +1,145 @@
+"""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 hub, nn
+
+from eva.core.utils.progress_bar import tqdm
+
+
+def load_model_weights(model: nn.Module, checkpoint_path: str) -> None:
+    """Loads (local or remote) weights to the model in-place.
+
+    Args:
+        model: The model to load the weights to.
+        checkpoint_path: The path to the model weights/checkpoint.
+    """
+    logger.info(f"Loading '{model.__class__.__name__}' model from checkpoint '{checkpoint_path}'")
+
+    fs = cloud_io.get_filesystem(checkpoint_path)
+    with fs.open(checkpoint_path, "rb") as file:
+        checkpoint = cloud_io._load(file, map_location="cpu")  # type: ignore
+        if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
+            checkpoint = checkpoint["state_dict"]
+
+        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))

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/src/eva/core/trainers/_recorder.py

@@ -129,7 +129,10 @@ class SessionRecorder:
     def _save_config(self) -> None:
         """Saves the config yaml with resolved env placeholders to the output directory."""
         if self.config_path:
-            config = OmegaConf.load(self.config_path)
+            config_fs = cloud_io.get_filesystem(self.config_path)
+            with config_fs.open(self.config_path, "r") as config_file:
+                config = OmegaConf.load(config_file)  # type: ignore
+
             fs = cloud_io.get_filesystem(self._output_dir, anon=False)
             with fs.open(os.path.join(self._output_dir, self._config_file), "w") as file:
                 config_yaml = OmegaConf.to_yaml(config, resolve=True)

kaiko_eva-0.2.2/src/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

kaiko_eva-0.2.2/src/eva/vision/data/__init__.py

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

kaiko_eva-0.2.2/src/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"]

kaiko_eva-0.2.2/src/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"]

kaiko_eva-0.2.2/src/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

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/src/eva/vision/data/datasets/__init__.py

@@ -16,9 +16,9 @@ from eva.vision.data.datasets.classification import (
 )
 from eva.vision.data.datasets.segmentation import (
     BCSS,
+    BTCV,
     CoNSeP,
     EmbeddingsSegmentationDataset,
-    ImageSegmentation,
     LiTS,
     LiTSBalanced,
     MoNuSAC,
@@ -29,6 +29,7 @@ from eva.vision.data.datasets.wsi import MultiWsiDataset, WsiDataset
 
 __all__ = [
     "BACH",
+    "BTCV",
     "BCSS",
     "BreaKHis",
     "BRACS",
@@ -43,7 +44,6 @@ __all__ = [
     "WsiClassificationDataset",
     "CoNSeP",
     "EmbeddingsSegmentationDataset",
-    "ImageSegmentation",
     "LiTS",
     "LiTSBalanced",
     "MoNuSAC",

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/src/eva/vision/data/datasets/classification/bach.py

@@ -8,12 +8,11 @@ from torchvision import tv_tensors
 from torchvision.datasets import folder, utils
 from typing_extensions import override
 
-from eva.vision.data.datasets import _utils, _validators, structs
-from eva.vision.data.datasets.classification import base
+from eva.vision.data.datasets import _utils, _validators, structs, vision
 from eva.vision.utils import io
 
 
-class BACH(base.ImageClassification):
+class BACH(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     """Dataset class for BACH images and corresponding targets."""
 
     _train_index_ranges: List[Tuple[int, int]] = [
@@ -125,7 +124,7 @@ class BACH(base.ImageClassification):
         )
 
     @override
-    def load_image(self, index: int) -> tv_tensors.Image:
+    def load_data(self, index: int) -> tv_tensors.Image:
         image_path, _ = self._samples[self._indices[index]]
         return io.read_image_as_tensor(image_path)
 

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/src/eva/vision/data/datasets/classification/bracs.py

@@ -8,12 +8,11 @@ from torchvision import tv_tensors
 from torchvision.datasets import folder
 from typing_extensions import override
 
-from eva.vision.data.datasets import _validators
-from eva.vision.data.datasets.classification import base
+from eva.vision.data.datasets import _validators, vision
 from eva.vision.utils import io
 
 
-class BRACS(base.ImageClassification):
+class BRACS(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     """Dataset class for BRACS images and corresponding targets."""
 
     _expected_dataset_lengths: Dict[str, int] = {
@@ -80,7 +79,7 @@ class BRACS(base.ImageClassification):
         )
 
     @override
-    def load_image(self, index: int) -> tv_tensors.Image:
+    def load_data(self, index: int) -> tv_tensors.Image:
         image_path, _ = self._samples[index]
         return io.read_image_as_tensor(image_path)
 

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/src/eva/vision/data/datasets/classification/breakhis.py

@@ -10,12 +10,11 @@ from torchvision import tv_tensors
 from torchvision.datasets import utils
 from typing_extensions import override
 
-from eva.vision.data.datasets import _validators, structs
-from eva.vision.data.datasets.classification import base
+from eva.vision.data.datasets import _validators, structs, vision
 from eva.vision.utils import io
 
 
-class BreaKHis(base.ImageClassification):
+class BreaKHis(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     """Dataset class for BreaKHis images and corresponding targets."""
 
     _resources: List[structs.DownloadResource] = [
@@ -145,7 +144,7 @@ class BreaKHis(base.ImageClassification):
         )
 
     @override
-    def load_image(self, index: int) -> tv_tensors.Image:
+    def load_data(self, index: int) -> tv_tensors.Image:
         image_path = self._image_files[self._indices[index]]
         return io.read_image_as_tensor(image_path)
 

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/src/eva/vision/data/datasets/classification/camelyon16.py

@@ -11,12 +11,11 @@ from torchvision import tv_tensors
 from torchvision.transforms.v2 import functional
 from typing_extensions import override
 
-from eva.vision.data.datasets import _validators, wsi
-from eva.vision.data.datasets.classification import base
+from eva.vision.data.datasets import _validators, vision, wsi
 from eva.vision.data.wsi.patching import samplers
 
 
-class Camelyon16(wsi.MultiWsiDataset, base.ImageClassification):
+class Camelyon16(wsi.MultiWsiDataset, vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     """Dataset class for Camelyon16 images and corresponding targets."""
 
     _val_slides = [
@@ -195,10 +194,10 @@ class Camelyon16(wsi.MultiWsiDataset, base.ImageClassification):
 
     @override
     def __getitem__(self, index: int) -> Tuple[tv_tensors.Image, torch.Tensor, Dict[str, Any]]:
-        return base.ImageClassification.__getitem__(self, index)
+        return vision.VisionDataset.__getitem__(self, index)
 
     @override
-    def load_image(self, index: int) -> tv_tensors.Image:
+    def load_data(self, index: int) -> tv_tensors.Image:
         image_array = wsi.MultiWsiDataset.__getitem__(self, index)
         return functional.to_image(image_array)
 

{kaiko_eva-0.2.0 → kaiko_eva-0.2.2}/src/eva/vision/data/datasets/classification/crc.py

@@ -8,12 +8,11 @@ from torchvision import tv_tensors
 from torchvision.datasets import folder, utils
 from typing_extensions import override
 
-from eva.vision.data.datasets import _validators, structs
-from eva.vision.data.datasets.classification import base
+from eva.vision.data.datasets import _validators, structs, vision
 from eva.vision.utils import io
 
 
-class CRC(base.ImageClassification):
+class CRC(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     """Dataset class for CRC images and corresponding targets."""
 
     _train_resource: structs.DownloadResource = structs.DownloadResource(
@@ -117,7 +116,7 @@ class CRC(base.ImageClassification):
         )
 
     @override
-    def load_image(self, index: int) -> tv_tensors.Image:
+    def load_data(self, index: int) -> tv_tensors.Image:
         image_path, _ = self._samples[index]
         return io.read_image_as_tensor(image_path)