kaiko-eva 0.2.0__py3-none-any.whl → 0.2.2__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/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)
 

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/_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)

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

@@ -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",

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)
 

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)
 

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)
 

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)
 

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)
 

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

@@ -12,12 +12,11 @@ from loguru import logger
 from torchvision import tv_tensors
 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 GleasonArvaniti(base.ImageClassification):
+class GleasonArvaniti(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     """Dataset class for GleasonArvaniti images and corresponding targets."""
 
     _expected_dataset_lengths: Dict[str | None, int] = {
@@ -121,7 +120,7 @@ class GleasonArvaniti(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)
 

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

@@ -7,12 +7,11 @@ import torch
 from torchvision import tv_tensors
 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 MHIST(base.ImageClassification):
+class MHIST(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     """MHIST dataset."""
 
     def __init__(
@@ -69,7 +68,7 @@ class MHIST(base.ImageClassification):
         )
 
     @override
-    def load_image(self, index: int) -> tv_tensors.Image:
+    def load_data(self, index: int) -> tv_tensors.Image:
         image_filename, _ = self._samples[index]
         image_path = os.path.join(self._dataset_path, image_filename)
         return io.read_image_as_tensor(image_path)

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

@@ -13,12 +13,11 @@ from torchvision.transforms.v2 import functional
 from typing_extensions import override
 
 from eva.core.data import splitting
-from eva.vision.data.datasets import _validators, structs, wsi
-from eva.vision.data.datasets.classification import base
+from eva.vision.data.datasets import _validators, structs, vision, wsi
 from eva.vision.data.wsi.patching import samplers
 
 
-class PANDA(wsi.MultiWsiDataset, base.ImageClassification):
+class PANDA(wsi.MultiWsiDataset, vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     """Dataset class for PANDA images and corresponding targets."""
 
     _train_split_ratio: float = 0.7
@@ -121,10 +120,10 @@ class PANDA(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)
 

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

@@ -10,14 +10,13 @@ from torchvision.datasets import utils
 from torchvision.transforms.v2 import functional
 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
 
 _URL_TEMPLATE = "https://zenodo.org/records/2546921/files/{filename}.gz?download=1"
 """PatchCamelyon URL files templates."""
 
 
-class PatchCamelyon(base.ImageClassification):
+class PatchCamelyon(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     """Dataset class for PatchCamelyon images and corresponding targets."""
 
     _train_resources: List[structs.DownloadResource] = [
@@ -127,7 +126,7 @@ class PatchCamelyon(base.ImageClassification):
         )
 
     @override
-    def load_image(self, index: int) -> tv_tensors.Image:
+    def load_data(self, index: int) -> tv_tensors.Image:
         return self._load_from_h5("x", index)
 
     @override

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

@@ -10,12 +10,11 @@ import torch
 from torchvision import tv_tensors
 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 UniToPatho(base.ImageClassification):
+class UniToPatho(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     """Dataset class for UniToPatho images and corresponding targets."""
 
     _expected_dataset_lengths: Dict[str | None, int] = {
@@ -109,7 +108,7 @@ class UniToPatho(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)
 

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

@@ -9,12 +9,13 @@ import torch
 from torchvision import tv_tensors
 from typing_extensions import override
 
-from eva.vision.data.datasets import wsi
-from eva.vision.data.datasets.classification import base
+from eva.vision.data.datasets import vision, wsi
 from eva.vision.data.wsi.patching import samplers
 
 
-class WsiClassificationDataset(wsi.MultiWsiDataset, base.ImageClassification):
+class WsiClassificationDataset(
+    wsi.MultiWsiDataset, vision.VisionDataset[tv_tensors.Image, torch.Tensor]
+):
     """A general dataset class for whole-slide image classification using manifest files."""
 
     default_column_mapping: Dict[str, str] = {
@@ -78,10 +79,10 @@ class WsiClassificationDataset(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:
         return wsi.MultiWsiDataset.__getitem__(self, index)
 
     @override