kaiko-eva 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

eva/language/models/wrappers/litellm.py

@@ -16,7 +16,7 @@ from litellm.exceptions import (
 from loguru import logger
 from typing_extensions import override
 
-from eva.language.models.typings import TextBatch
+from eva.language.models.typings import ModelOutput, TextBatch
 from eva.language.models.wrappers import base
 from eva.language.utils.text import messages as message_utils
 
@@ -32,6 +32,14 @@ RETRYABLE_ERRORS = (
 class LiteLLMModel(base.LanguageModel):
     """Wrapper class for LiteLLM language models."""
 
+    _default_model_kwargs = {
+        "temperature": 0.0,
+        "max_completion_tokens": 1024,
+        "top_p": 1.0,
+        "seed": 42,
+    }
+    """Default API model parameters for evaluation."""
+
     def __init__(
         self,
         model_name: str,
@@ -51,9 +59,10 @@ class LiteLLMModel(base.LanguageModel):
         super().__init__(system_prompt=system_prompt)
 
         self.model_name = model_name
-        self.model_kwargs = model_kwargs or {}
+        self.model_kwargs = self._default_model_kwargs | (model_kwargs or {})
 
         litellm.suppress_debug_info = True
+        litellm.drop_params = True
 
         if log_level is not None:
             logging.getLogger("LiteLLM").setLevel(log_level)
@@ -94,16 +103,17 @@ class LiteLLMModel(base.LanguageModel):
             f"Retrying due to {details.get('exception') or 'Unknown error'}"
         ),
     )
-    def model_forward(self, batch: List[List[Dict[str, Any]]]) -> List[str]:
+    def model_forward(self, batch: List[List[Dict[str, Any]]]) -> ModelOutput:
         """Generates output text through API calls via LiteLLM's batch completion functionality."""
         outputs = batch_completion(model=self.model_name, messages=batch, **self.model_kwargs)
         self._raise_exceptions(outputs)
 
-        return [
+        generated_text = [
             output["choices"][0]["message"]["content"]
             for output in outputs
             if output["choices"][0]["message"]["role"] == "assistant"
         ]
+        return ModelOutput(generated_text=generated_text)
 
     def _raise_exceptions(self, outputs: list):
         for output in outputs:

eva/multimodal/models/modules/vision_language.py

@@ -1,6 +1,6 @@
 """Model module for vision-language models."""
 
-from typing import Any, List
+from typing import Any
 
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from torch import nn
@@ -9,6 +9,7 @@ from typing_extensions import override
 from eva.core.metrics import structs as metrics_lib
 from eva.core.models.modules import module
 from eva.core.models.modules.utils import batch_postprocess
+from eva.language.models.typings import ModelOutput
 from eva.multimodal.models.typings import TextImageBatch
 
 
@@ -33,7 +34,7 @@ class VisionLanguageModule(module.ModelModule):
         self.model = model
 
     @override
-    def forward(self, batch: TextImageBatch, *args: Any, **kwargs: Any) -> List[str]:
+    def forward(self, batch: TextImageBatch, *args: Any, **kwargs: Any) -> ModelOutput:
         return self.model(batch)
 
     @override
@@ -46,10 +47,10 @@ class VisionLanguageModule(module.ModelModule):
 
     def _batch_step(self, batch: TextImageBatch) -> STEP_OUTPUT:
         text, _, targets, metadata = TextImageBatch(*batch)
-        predictions = self.forward(batch)
+        output = self.forward(batch)
         return {
             "inputs": text,
-            "predictions": predictions,
+            "predictions": output.pop("generated_text"),  # type: ignore
             "targets": targets,
             "metadata": metadata,
-        }
+        } | output

eva/multimodal/models/networks/alibaba.py

@@ -36,4 +36,5 @@ class Qwen25VL7BInstruct(wrappers.HuggingFaceModel):
                 "max_pixels": 451584,  # 672*672
             },
             system_prompt=system_prompt,
+            image_key="images",
         )

eva/multimodal/models/networks/others.py

@@ -20,7 +20,7 @@ class PathoR13b(wrappers.HuggingFaceModel):
         attn_implementation: str = "flash_attention_2",
     ):
         """Initialize the Patho-R1-3B model."""
-        requirements.check_dependencies(requirements={"torch": "2.5.1", "torchvision": "0.20.1"})
+        requirements.check_min_versions(requirements={"torch": "2.5.1", "torchvision": "0.20.1"})
 
         if not os.getenv("HF_TOKEN"):
             raise ValueError("HF_TOKEN env variable must be set.")
@@ -44,4 +44,5 @@ class PathoR13b(wrappers.HuggingFaceModel):
                 "max_pixels": 451584,  # 672*672
             },
             system_prompt=system_prompt,
+            image_key="images",
         )

eva/multimodal/models/wrappers/base.py

@@ -1,16 +1,17 @@
 """Base class for vision language model wrappers."""
 
 import abc
-from typing import Any, Callable, List
+from typing import Any, Callable
 
 from typing_extensions import override
 
 from eva.core.models.wrappers import base
 from eva.language.data.messages import ModelSystemMessage
+from eva.language.models.typings import ModelOutput
 from eva.multimodal.models.typings import TextImageBatch
 
 
-class VisionLanguageModel(base.BaseModel[TextImageBatch, List[str]]):
+class VisionLanguageModel(base.BaseModel[TextImageBatch, ModelOutput]):
     """Base class for multimodal models.
 
     Classes that inherit from this should implement the following methods:
@@ -36,7 +37,7 @@ class VisionLanguageModel(base.BaseModel[TextImageBatch, List[str]]):
         self.system_message = ModelSystemMessage(content=system_prompt) if system_prompt else None
 
     @override
-    def forward(self, batch: TextImageBatch) -> List[str]:
+    def forward(self, batch: TextImageBatch) -> ModelOutput:
         """Forward pass of the model."""
         inputs = self.format_inputs(batch)
         return super().forward(inputs)

eva/multimodal/models/wrappers/huggingface.py

@@ -9,10 +9,11 @@ from loguru import logger
 from torch import nn
 from typing_extensions import override
 
-from eva.language.models.typings import TextBatch
+from eva.language.models.typings import ModelOutput, TextBatch
 from eva.language.utils.text import messages as language_message_utils
 from eva.multimodal.models.typings import TextImageBatch
 from eva.multimodal.models.wrappers import base
+from eva.multimodal.utils.batch import unpack_batch
 from eva.multimodal.utils.text import messages as message_utils
 
 
@@ -27,6 +28,14 @@ class HuggingFaceModel(base.VisionLanguageModel):
         generation_kwargs: Additional generation arguments.
     """
 
+    _default_generation_kwargs = {
+        "temperature": 0.0,
+        "max_new_tokens": 1024,
+        "do_sample": False,
+        "top_p": 1.0,
+    }
+    """Default HF model parameters for evaluation."""
+
     def __init__(
         self,
         model_name_or_path: str,
@@ -35,6 +44,7 @@ class HuggingFaceModel(base.VisionLanguageModel):
         system_prompt: str | None = None,
         processor_kwargs: Dict[str, Any] | None = None,
         generation_kwargs: Dict[str, Any] | None = None,
+        image_key: str = "image",
     ):
         """Initialize the HuggingFace model wrapper.
 
@@ -45,6 +55,7 @@ class HuggingFaceModel(base.VisionLanguageModel):
             system_prompt: System prompt to use.
             processor_kwargs: Additional processor arguments.
             generation_kwargs: Additional generation arguments.
+            image_key: The key used for image inputs in the chat template.
         """
         super().__init__(system_prompt=system_prompt)
 
@@ -52,7 +63,8 @@ class HuggingFaceModel(base.VisionLanguageModel):
         self.model_kwargs = model_kwargs or {}
         self.base_model_class = model_class
         self.processor_kwargs = processor_kwargs or {}
-        self.generation_kwargs = generation_kwargs or {}
+        self.generation_kwargs = self._default_generation_kwargs | (generation_kwargs or {})
+        self.image_key = image_key
 
         self.processor = self.load_processor()
         self.model = self.load_model()
@@ -72,7 +84,7 @@ class HuggingFaceModel(base.VisionLanguageModel):
                 "pixel_values": ...
             }
         """
-        message_batch, image_batch, _, _ = self._unpack_batch(batch)
+        message_batch, image_batch, _, _ = unpack_batch(batch)
         with_images = image_batch is not None
 
         message_batch = language_message_utils.batch_insert_system_message(
@@ -105,12 +117,12 @@ class HuggingFaceModel(base.VisionLanguageModel):
         }
 
         if with_images:
-            processor_inputs["image"] = [[image] for image in image_batch]
+            processor_inputs[self.image_key] = [[image] for image in image_batch]
 
         return self.processor(**processor_inputs).to(self.model.device)  # type: ignore
 
     @override
-    def model_forward(self, batch: Dict[str, torch.Tensor]) -> List[str]:
+    def model_forward(self, batch: Dict[str, torch.Tensor]) -> ModelOutput:
         """Generates text output from the model. Is called by the `generate` method.
 
         Args:
@@ -121,8 +133,14 @@ class HuggingFaceModel(base.VisionLanguageModel):
         Returns:
             A dictionary containing the processed input and the model's output.
         """
-        output = self.model.generate(**batch, **self.generation_kwargs)  # type: ignore
-        return self._decode_output(output, batch["input_ids"].shape[-1])
+        output_ids = self.model.generate(**batch, **self.generation_kwargs)  # type: ignore
+
+        return ModelOutput(
+            generated_text=self._decode_output(output_ids, batch["input_ids"].shape[-1]),
+            input_ids=batch.get("input_ids"),
+            output_ids=output_ids,
+            attention_mask=batch.get("attention_mask"),
+        )
 
     @override
     def load_model(self) -> nn.Module:
@@ -148,15 +166,10 @@ class HuggingFaceModel(base.VisionLanguageModel):
     def load_processor(self) -> Callable:
         """Initialize the processor."""
         return transformers.AutoProcessor.from_pretrained(
-            self.model_name_or_path,
+            self.processor_kwargs.pop("model_name_or_path", self.model_name_or_path),
             **self.processor_kwargs,
         )
 
-    def _unpack_batch(self, batch: TextImageBatch | TextBatch) -> tuple:
-        if isinstance(batch, TextImageBatch):
-            return batch.text, batch.image, batch.target, batch.metadata
-        return batch.text, None, batch.target, batch.metadata
-
     def _decode_output(self, output: torch.Tensor, instruction_length: int) -> List[str]:
         """Decode the model's batch output to text.
 

eva/multimodal/models/wrappers/litellm.py

@@ -6,9 +6,11 @@ from typing import Any, Dict, List
 from typing_extensions import override
 
 from eva.language.models import wrappers as language_wrappers
+from eva.language.models.typings import ModelOutput
 from eva.language.utils.text import messages as language_message_utils
 from eva.multimodal.models.typings import TextImageBatch
 from eva.multimodal.models.wrappers import base
+from eva.multimodal.utils.batch import unpack_batch
 from eva.multimodal.utils.text import messages as message_utils
 
 
@@ -42,7 +44,7 @@ class LiteLLMModel(base.VisionLanguageModel):
 
     @override
     def format_inputs(self, batch: TextImageBatch) -> List[List[Dict[str, Any]]]:
-        message_batch, image_batch, _, _ = TextImageBatch(*batch)
+        message_batch, image_batch, _, _ = unpack_batch(batch)
 
         message_batch = language_message_utils.batch_insert_system_message(
             message_batch, self.system_message
@@ -52,5 +54,5 @@ class LiteLLMModel(base.VisionLanguageModel):
         return list(map(message_utils.format_litellm_message, message_batch, image_batch))
 
     @override
-    def model_forward(self, batch: List[List[Dict[str, Any]]]) -> List[str]:
+    def model_forward(self, batch: List[List[Dict[str, Any]]]) -> ModelOutput:
         return self.language_model.model_forward(batch)

eva/multimodal/utils/batch/__init__.py

@@ -0,0 +1,5 @@
+"""Multimodal batch utilities API."""
+
+from eva.multimodal.utils.batch.unpack import unpack_batch
+
+__all__ = ["unpack_batch"]

eva/multimodal/utils/batch/unpack.py

@@ -0,0 +1,11 @@
+"""Unpack batch utility function."""
+
+from eva.language.models.typings import TextBatch
+from eva.multimodal.models.typings import TextImageBatch
+
+
+def unpack_batch(batch: TextImageBatch | TextBatch) -> tuple:
+    """Unpacks a TextImageBatch or TextBatch into its components."""
+    if isinstance(batch, TextImageBatch):
+        return batch.text, batch.image, batch.target, batch.metadata
+    return batch.text, None, batch.target, batch.metadata

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

@@ -101,11 +101,6 @@ class BreaKHis(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     def class_to_idx(self) -> Dict[str, int]:
         return {label: index for index, label in enumerate(self.classes)}
 
-    @property
-    def _dataset_path(self) -> str:
-        """Returns the path of the image data of the dataset."""
-        return os.path.join(self._root, "BreaKHis_v1", "histology_slides")
-
     @functools.cached_property
     def _image_files(self) -> List[str]:
         """Return the list of image files in the dataset.
@@ -115,14 +110,14 @@ class BreaKHis(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
         """
         image_files = []
         for magnification in self._magnifications:
-            files_pattern = os.path.join(self._dataset_path, f"**/{magnification}", "*.png")
+            files_pattern = os.path.join(self._root, f"**/{magnification}", "*.png")
             image_files.extend(list(glob.glob(files_pattern, recursive=True)))
         return sorted(image_files)
 
     @override
     def filename(self, index: int) -> str:
         image_path = self._image_files[self._indices[index]]
-        return os.path.relpath(image_path, self._dataset_path)
+        return os.path.relpath(image_path, self._root)
 
     @override
     def prepare_data(self) -> None:
@@ -136,6 +131,8 @@ class BreaKHis(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
 
     @override
     def validate(self) -> None:
+        if not os.path.exists(self._root):
+            raise RuntimeError(f"Dataset not found at {self._root}.")
         _validators.check_dataset_integrity(
             self,
             length=self._expected_dataset_lengths[self._split],
@@ -164,7 +161,7 @@ class BreaKHis(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     def _download_dataset(self) -> None:
         """Downloads the dataset."""
         for resource in self._resources:
-            if os.path.isdir(self._dataset_path):
+            if os.path.isdir(self._root):
                 continue
 
             self._print_license()

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

@@ -12,6 +12,7 @@ from torchvision.datasets import utils
 from torchvision.transforms.v2 import functional
 from typing_extensions import override
 
+from eva.core import utils as core_utils
 from eva.core.data import splitting
 from eva.vision.data.datasets import _validators, structs, vision, wsi
 from eva.vision.data.wsi.patching import samplers
@@ -50,6 +51,7 @@ class PANDA(wsi.MultiWsiDataset, vision.VisionDataset[tv_tensors.Image, torch.Te
         image_transforms: Callable | None = None,
         coords_path: str | None = None,
         seed: int = 42,
+        download_dir: str | None = None,
     ) -> None:
         """Initializes the dataset.
 
@@ -64,10 +66,13 @@ class PANDA(wsi.MultiWsiDataset, vision.VisionDataset[tv_tensors.Image, torch.Te
             image_transforms: Transforms to apply to the extracted image patches.
             coords_path: File path to save the patch coordinates as .csv.
             seed: Random seed for reproducibility.
+            download_dir: Directory to download the dataset resources to. If None,
+                defaults to eva's home directory.
         """
         self._split = split
         self._root = root
         self._seed = seed
+        self._download_dir = download_dir or os.path.join(core_utils.home_dir(), "data", "panda")
 
         self._download_resources()
 
@@ -92,7 +97,7 @@ class PANDA(wsi.MultiWsiDataset, vision.VisionDataset[tv_tensors.Image, torch.Te
     @functools.cached_property
     def annotations(self) -> pd.DataFrame:
         """Loads the dataset labels."""
-        path = os.path.join(self._root, "train_with_noisy_labels.csv")
+        path = os.path.join(self._download_dir, "train_with_noisy_labels.csv")
         return pd.read_csv(path, index_col="image_id")
 
     @override
@@ -100,14 +105,16 @@ class PANDA(wsi.MultiWsiDataset, vision.VisionDataset[tv_tensors.Image, torch.Te
         _validators.check_dataset_exists(self._root, False)
 
         if not os.path.isdir(os.path.join(self._root, "train_images")):
-            raise FileNotFoundError("'train_images' directory not found in the root folder.")
-        if not os.path.isfile(os.path.join(self._root, "train_with_noisy_labels.csv")):
-            raise FileNotFoundError("'train.csv' file not found in the root folder.")
+            raise FileNotFoundError(f"'train_images' dir not found in folder: {self._root}")
+        if not os.path.isfile(os.path.join(self._download_dir, "train_with_noisy_labels.csv")):
+            raise FileNotFoundError(
+                f"'train_with_noisy_labels.csv' file not found in folder: {self._download_dir}"
+            )
 
     def _download_resources(self) -> None:
         """Downloads the dataset resources."""
         for resource in self._resources:
-            utils.download_url(resource.url, self._root, resource.filename, resource.md5)
+            utils.download_url(resource.url, self._download_dir, resource.filename, resource.md5)
 
     @override
     def validate(self) -> None:

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

@@ -106,7 +106,7 @@ class BTCV(VisionDataset[eva_tv_tensors.Volume, tv_tensors.Mask]):
 
     @override
     def validate(self) -> None:
-        requirements.check_dependencies(requirements={"torch": "2.5.1", "torchvision": "0.20.1"})
+        requirements.check_min_versions(requirements={"torch": "2.5.1", "torchvision": "0.20.1"})
 
         def _valid_sample(index: int) -> bool:
             """Indicates if the sample files exist and are reachable."""

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

@@ -108,7 +108,7 @@ class CoNSeP(wsi.MultiWsiDataset, vision.VisionDataset[tv_tensors.Image, tv_tens
             n_classes=5,
             first_and_last_labels=((self.classes[0], self.classes[-1])),
         )
-        n_expected = self._expected_dataset_lengths[None]
+        n_expected = self._expected_dataset_lengths[self._split]
         if len(self._file_paths) != n_expected:
             raise ValueError(
                 f"Expected {n_expected} images, found {len(self._file_paths)} in {self._root}."

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

@@ -123,7 +123,7 @@ class LiTS17(VisionDataset[eva_tv_tensors.Volume, tv_tensors.Mask]):
 
     @override
     def validate(self) -> None:
-        requirements.check_dependencies(requirements={"torch": "2.5.1", "torchvision": "0.20.1"})
+        requirements.check_min_versions(requirements={"torch": "2.5.1", "torchvision": "0.20.1"})
 
         def _valid_sample(index: int) -> bool:
             """Indicates if the sample files exist and are reachable."""

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

@@ -15,6 +15,7 @@ from torchvision import tv_tensors
 from torchvision.datasets import utils
 from typing_extensions import override
 
+from eva.core import utils as core_utils
 from eva.core.utils.progress_bar import tqdm
 from eva.vision.data.datasets import _validators, structs, vision
 from eva.vision.utils import io
@@ -55,6 +56,7 @@ class MoNuSAC(vision.VisionDataset[tv_tensors.Image, tv_tensors.Mask]):
         root: str,
         split: Literal["train", "test"],
         export_masks: bool = True,
+        processed_dir: str | None = None,
         download: bool = False,
         transforms: Callable | None = None,
     ) -> None:
@@ -66,6 +68,8 @@ class MoNuSAC(vision.VisionDataset[tv_tensors.Image, tv_tensors.Mask]):
             split: Dataset split to use.
             export_masks: Whether to export, save and use the semantic label masks
                 from disk.
+            processed_dir: Directory where to store the processed masks.
+                Only used if `export_masks` is `True`.
             download: Whether to download the data for the specified split.
                 Note that the download will be executed only by additionally
                 calling the :meth:`prepare_data` method and if the data does not
@@ -79,6 +83,9 @@ class MoNuSAC(vision.VisionDataset[tv_tensors.Image, tv_tensors.Mask]):
         self._split = split
         self._export_masks = export_masks
         self._download = download
+        self._processed_dir = processed_dir or os.path.join(
+            core_utils.home_dir(), "data", "processed", "monusac"
+        )
 
     @property
     @override
@@ -155,10 +162,7 @@ class MoNuSAC(vision.VisionDataset[tv_tensors.Image, tv_tensors.Mask]):
 
     def _export_semantic_label_masks(self) -> None:
         """Export semantic label masks to disk."""
-        mask_files = [
-            (index, filename.replace(".tif", ".npy"))
-            for index, filename in enumerate(self._image_files)
-        ]
+        mask_files = [(i, self._processed_filename(i)) for i in range(len(self._image_files))]
         to_export = filter(lambda x: not os.path.isfile(x[1]), mask_files)
         for sample_index, filename in tqdm(
             list(to_export),
@@ -166,6 +170,7 @@ class MoNuSAC(vision.VisionDataset[tv_tensors.Image, tv_tensors.Mask]):
             leave=False,
         ):
             semantic_labels = self._get_semantic_mask(sample_index)
+            os.makedirs(os.path.dirname(filename), exist_ok=True)
             np.save(filename, semantic_labels)
 
     def _load_semantic_mask_file(self, index: int) -> npt.NDArray[Any]:
@@ -177,8 +182,7 @@ class MoNuSAC(vision.VisionDataset[tv_tensors.Image, tv_tensors.Mask]):
         Returns:
             Loaded mask as a numpy array.
         """
-        mask_filename = self._image_files[index].replace(".tif", ".npy")
-        return np.load(mask_filename)
+        return np.load(self._processed_filename(index))
 
     def _get_semantic_mask(self, index: int) -> npt.NDArray[Any]:
         """Builds and loads the semantic label mask from the XML annotations.
@@ -216,6 +220,11 @@ class MoNuSAC(vision.VisionDataset[tv_tensors.Image, tv_tensors.Mask]):
 
         return semantic_labels
 
+    def _processed_filename(self, index: int) -> str:
+        """Returns the path of the processed mask for a given index."""
+        relative_path = os.path.relpath(self._image_files[index], self._root)
+        return os.path.join(self._processed_dir, relative_path).replace(".tif", ".npy")
+
     def _download_dataset(self) -> None:
         """Downloads the dataset."""
         self._print_license()

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

@@ -95,7 +95,7 @@ class MSDTask7Pancreas(VisionDataset[eva_tv_tensors.Volume, tv_tensors.Mask]):
 
     @override
     def validate(self) -> None:
-        requirements.check_dependencies(requirements={"torch": "2.5.1", "torchvision": "0.20.1"})
+        requirements.check_min_versions(requirements={"torch": "2.5.1", "torchvision": "0.20.1"})
 
         def _valid_sample(index: int) -> bool:
             """Indicates if the sample files exist and are reachable."""

eva/vision/data/transforms/base/__init__.py

@@ -1,5 +1,6 @@
 """Base classes for transforms."""
 
 from eva.vision.data.transforms.base.monai import RandomMonaiTransform
+from eva.vision.data.transforms.base.torchvision import TorchvisionTransformV2
 
-__all__ = ["RandomMonaiTransform"]
+__all__ = ["RandomMonaiTransform", "TorchvisionTransformV2"]

eva/vision/data/transforms/base/monai.py

@@ -2,10 +2,10 @@
 
 import abc
 
-from torchvision.transforms import v2
+from eva.vision.data.transforms.base.torchvision import TorchvisionTransformV2
 
 
-class RandomMonaiTransform(v2.Transform, abc.ABC):
+class RandomMonaiTransform(TorchvisionTransformV2, abc.ABC):
     """Base class for MONAI transform wrappers."""
 
     @abc.abstractmethod

eva/vision/data/transforms/base/torchvision.py

@@ -0,0 +1,33 @@
+"""Base class for torchvision.v2 transforms."""
+
+import abc
+from typing import Any, Dict, List
+
+from torchvision.transforms import v2
+
+
+class TorchvisionTransformV2(v2.Transform, abc.ABC):
+    """Wrapper for torchvision.v2.Transform.
+
+    This class ensures compatibility both with >=0.21.0 and older versions,
+    as torchvision 0.21.0 introduced a new transform API where they
+    renamed the following methods:
+
+    - `_get_params` -> `make_params`
+    - `_transform` -> `transform`
+    """
+
+    def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+        """Called internally before calling transform() on each input."""
+        return {}
+
+    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+        return self.make_params(flat_inputs)
+
+    @abc.abstractmethod
+    def transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        """Applies the transformation to the input."""
+        raise NotImplementedError
+
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return self.transform(inpt, params)

eva/vision/data/transforms/common/squeeze.py

@@ -4,10 +4,12 @@ from typing import Any
 
 import torch
 from torchvision import tv_tensors
-from torchvision.transforms import v2
+from typing_extensions import override
 
+from eva.vision.data.transforms import base
 
-class Squeeze(v2.Transform):
+
+class Squeeze(base.TorchvisionTransformV2):
     """Squeezes the input tensor accross all or specified dimensions."""
 
     def __init__(self, dim: int | list[int] | None = None):
@@ -19,6 +21,7 @@ class Squeeze(v2.Transform):
         super().__init__()
         self._dim = dim
 
-    def _transform(self, inpt: Any, params: dict[str, Any]) -> Any:
+    @override
+    def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
         output = torch.squeeze(inpt) if self._dim is None else torch.squeeze(inpt, dim=self._dim)
         return tv_tensors.wrap(output, like=inpt)

eva/vision/data/transforms/croppad/crop_foreground.py

@@ -8,13 +8,13 @@ from monai.config import type_definitions
 from monai.transforms.croppad import array as monai_croppad_transforms
 from monai.utils.enums import PytorchPadMode
 from torchvision import tv_tensors
-from torchvision.transforms import v2
 from typing_extensions import override
 
 from eva.vision.data import tv_tensors as eva_tv_tensors
+from eva.vision.data.transforms import base
 
 
-class CropForeground(v2.Transform):
+class CropForeground(base.TorchvisionTransformV2):
     """Crop an image using a bounding box.
 
     The bounding box is generated by selecting foreground using select_fn
@@ -74,19 +74,20 @@ class CropForeground(v2.Transform):
             **pad_kwargs,
         )
 
-    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+    @override
+    def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
         volume = next(inpt for inpt in flat_inputs if isinstance(inpt, eva_tv_tensors.Volume))
         box_start, box_end = self._foreground_crop.compute_bounding_box(volume)
         return {"box_start": box_start, "box_end": box_end}
 
     @functools.singledispatchmethod
     @override
-    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+    def transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
         return inpt
 
-    @_transform.register(tv_tensors.Image)
-    @_transform.register(eva_tv_tensors.Volume)
-    @_transform.register(tv_tensors.Mask)
+    @transform.register(tv_tensors.Image)
+    @transform.register(eva_tv_tensors.Volume)
+    @transform.register(tv_tensors.Mask)
     def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
         inpt_foreground_cropped = self._foreground_crop.crop_pad(
             inpt, params["box_start"], params["box_end"]