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

eva/multimodal/models/wrappers/huggingface.py

@@ -0,0 +1,193 @@
+"""HuggingFace Vision-Language Model Wrapper."""
+
+import functools
+from typing import Any, Callable, Dict, List
+
+import torch
+import transformers
+from loguru import logger
+from torch import nn
+from typing_extensions import override
+
+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
+
+
+class HuggingFaceModel(base.VisionLanguageModel):
+    """Lightweight wrapper for Huggingface VLMs.
+
+    Args:
+        model_name_or_path: The name of the model to use.
+        model_class: The class of the model to use.
+        model_kwargs: Additional model arguments.
+        processor_kwargs: Additional processor arguments.
+        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,
+        model_class: str,
+        model_kwargs: Dict[str, Any] | None = None,
+        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.
+
+        Args:
+            model_name_or_path: The name or path of the model to use.
+            model_class: The class of the model to use.
+            model_kwargs: Additional model arguments.
+            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)
+
+        self.model_name_or_path = model_name_or_path
+        self.model_kwargs = model_kwargs or {}
+        self.base_model_class = model_class
+        self.processor_kwargs = processor_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()
+
+    @override
+    def format_inputs(self, batch: TextImageBatch | TextBatch) -> Dict[str, torch.Tensor]:
+        """Formats inputs for HuggingFace models.
+
+        Args:
+            batch: A batch of text and image inputs.
+
+        Returns:
+            A dictionary produced by the provided processor following a format like:
+            {
+                "input_ids": ...,
+                "attention_mask": ...,
+                "pixel_values": ...
+            }
+        """
+        message_batch, image_batch, _, _ = unpack_batch(batch)
+        with_images = image_batch is not None
+
+        message_batch = language_message_utils.batch_insert_system_message(
+            message_batch, self.system_message
+        )
+        message_batch = list(map(language_message_utils.combine_system_messages, message_batch))
+
+        if self.processor.chat_template is not None:  # type: ignore
+            templated_text = [
+                self.processor.apply_chat_template(  # type: ignore
+                    message,
+                    add_generation_prompt=True,
+                    tokenize=False,
+                )
+                for message in map(
+                    functools.partial(
+                        message_utils.format_huggingface_message,
+                        with_images=with_images,
+                    ),
+                    message_batch,
+                )
+            ]
+        else:
+            raise NotImplementedError("Currently only chat models are supported.")
+
+        processor_inputs = {
+            "text": templated_text,
+            "return_tensors": "pt",
+            **self.processor_kwargs,
+        }
+
+        if with_images:
+            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]) -> ModelOutput:
+        """Generates text output from the model. Is called by the `generate` method.
+
+        Args:
+            batch: A dictionary containing the input data, which may include:
+                - "text": List of messages formatted for the model.
+                - "image": List of image tensors.
+
+        Returns:
+            A dictionary containing the processed input and the model's output.
+        """
+        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:
+        """Setting up the model. Used for delayed model initialization.
+
+        Raises:
+            ValueError: If the model class is not found in transformers or if the model
+                does not support gradient checkpointing but it is enabled.
+        """
+        logger.info(f"Configuring model: {self.model_name_or_path}")
+        if hasattr(transformers, self.base_model_class):
+            model_class = getattr(transformers, self.base_model_class)
+        else:
+            raise ValueError(f"Model class {self.base_model_class} not found in transformers")
+
+        model = model_class.from_pretrained(self.model_name_or_path, **self.model_kwargs)
+
+        if not hasattr(model, "generate"):
+            raise ValueError(f"Model {self.model_name_or_path} does not support generation. ")
+
+        return model
+
+    def load_processor(self) -> Callable:
+        """Initialize the processor."""
+        return transformers.AutoProcessor.from_pretrained(
+            self.processor_kwargs.pop("model_name_or_path", self.model_name_or_path),
+            **self.processor_kwargs,
+        )
+
+    def _decode_output(self, output: torch.Tensor, instruction_length: int) -> List[str]:
+        """Decode the model's batch output to text.
+
+        Args:
+            output: The raw output from the model.
+            instruction_length: The length of the instruction in the input.
+
+        Returns:
+            A list of decoded text responses.
+        """
+        decoded_input = self.processor.batch_decode(  # type: ignore
+            output[:, :instruction_length], skip_special_tokens=True
+        )
+        decoded_output = self.processor.batch_decode(  # type: ignore
+            output[:, instruction_length:], skip_special_tokens=True
+        )
+
+        logger.debug(f"Decoded input: {decoded_input}")
+        logger.debug(f"Decoded output: {decoded_output}")
+
+        return decoded_output

eva/multimodal/models/wrappers/litellm.py

@@ -0,0 +1,58 @@
+"""LiteLLM vision-language model wrapper."""
+
+import logging
+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
+
+
+class LiteLLMModel(base.VisionLanguageModel):
+    """Wrapper class for LiteLLM vision-language models."""
+
+    def __init__(
+        self,
+        model_name: str,
+        model_kwargs: Dict[str, Any] | None = None,
+        system_prompt: str | None = None,
+        log_level: int | None = logging.INFO,
+    ):
+        """Initialize the LiteLLM Wrapper.
+
+        Args:
+            model_name: The name of the model to use.
+            model_kwargs: Additional keyword arguments to pass during
+                generation (e.g., `temperature`, `max_tokens`).
+            system_prompt: The system prompt to use (optional).
+            log_level: Optional logging level for LiteLLM. Defaults to WARNING.
+        """
+        super().__init__(system_prompt=system_prompt)
+
+        self.language_model = language_wrappers.LiteLLMModel(
+            model_name=model_name,
+            model_kwargs=model_kwargs,
+            system_prompt=system_prompt,
+            log_level=log_level,
+        )
+
+    @override
+    def format_inputs(self, batch: TextImageBatch) -> List[List[Dict[str, Any]]]:
+        message_batch, image_batch, _, _ = unpack_batch(batch)
+
+        message_batch = language_message_utils.batch_insert_system_message(
+            message_batch, self.system_message
+        )
+        message_batch = list(map(language_message_utils.combine_system_messages, message_batch))
+
+        return list(map(message_utils.format_litellm_message, message_batch, image_batch))
+
+    @override
+    def model_forward(self, batch: List[List[Dict[str, Any]]]) -> ModelOutput:
+        return self.language_model.model_forward(batch)

eva/multimodal/utils/__init__.py

@@ -0,0 +1 @@
+"""Multimodal utilities API."""

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/multimodal/utils/image/__init__.py

@@ -0,0 +1,5 @@
+"""Multimodal image utilities API."""
+
+from eva.multimodal.utils.image.encode import encode_image
+
+__all__ = ["encode_image"]

eva/multimodal/utils/image/encode.py

@@ -0,0 +1,28 @@
+"""Image encoding utilities."""
+
+import base64
+import io
+from typing import Literal
+
+from torchvision import tv_tensors
+from torchvision.transforms.v2 import functional as F
+
+
+def encode_image(image: tv_tensors.Image, encoding: Literal["base64"]) -> str:
+    """Encodes an image tensor into a string format.
+
+    Args:
+        image: The image tensor to encode.
+        encoding: The encoding format to use. Currently only supports "base64".
+
+    Returns:
+        An encoded string representation of the image.
+    """
+    match encoding:
+        case "base64":
+            image_bytes = io.BytesIO()
+            F.to_pil_image(image).save(image_bytes, format="PNG", optimize=True)
+            image_bytes.seek(0)
+            return base64.b64encode(image_bytes.getvalue()).decode("utf-8")
+        case _:
+            raise ValueError(f"Unsupported encoding type: {encoding}. Supported: 'base64'")

eva/multimodal/utils/text/__init__.py

@@ -0,0 +1 @@
+"""Multimodal text utilities API."""

eva/multimodal/utils/text/messages.py

@@ -0,0 +1,79 @@
+"""Message formatting utilities for multimodal models."""
+
+from typing import Any, Dict, List
+
+from torchvision import tv_tensors
+
+from eva.language import utils as language_utils
+from eva.language.data.messages import MessageSeries, Role
+from eva.multimodal.utils import image as image_utils
+
+
+def format_huggingface_message(
+    message: MessageSeries, with_images: bool = False
+) -> List[Dict[str, Any]]:
+    """Formats a message series into a format suitable for Huggingface models."""
+    if not with_images:
+        return language_utils.format_chat_message(message)
+
+    formatted_message = []
+    for item in message:
+        if item.role == Role.SYSTEM:
+            formatted_message += language_utils.format_chat_message([item])
+        else:
+            formatted_message.append(
+                {
+                    "role": item.role,
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": str(item.content),
+                        },
+                        {"type": "image"},
+                    ],
+                }
+            )
+    return formatted_message
+
+
+def format_litellm_message(
+    message: MessageSeries, image: tv_tensors.Image | None
+) -> List[Dict[str, Any]]:
+    """Format a message series for LiteLLM API.
+
+    Args:
+        message: The message series to format.
+        image: Optional image to include in the message.
+
+    Returns:
+        A list of formatted message dictionaries.
+    """
+    if image is None:
+        return language_utils.format_chat_message(message)
+
+    formatted_message = []
+    for item in message:
+        if item.role == Role.SYSTEM:
+            formatted_message += language_utils.format_chat_message([item])
+        else:
+            formatted_message.append(
+                {
+                    "role": item.role,
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": str(item.content),
+                        },
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": (
+                                    f"data:image/png;base64,"
+                                    f"{image_utils.encode_image(image, encoding='base64')}"
+                                )
+                            },
+                        },
+                    ],
+                }
+            )
+    return formatted_message

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/classification/patch_camelyon.py

@@ -61,6 +61,13 @@ class PatchCamelyon(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
     ]
     """Test resources."""
 
+    _expected_length = {
+        "train": 262144,
+        "val": 32768,
+        "test": 32768,
+    }
+    """Expected dataset length for each split."""
+
     _license: str = (
         "Creative Commons Zero v1.0 Universal (https://choosealicense.com/licenses/cc0-1.0/)"
     )
@@ -113,14 +120,9 @@ class PatchCamelyon(vision.VisionDataset[tv_tensors.Image, torch.Tensor]):
 
     @override
     def validate(self) -> None:
-        expected_length = {
-            "train": 262144,
-            "val": 32768,
-            "test": 32768,
-        }
         _validators.check_dataset_integrity(
             self,
-            length=expected_length.get(self._split, 0),
+            length=self._expected_length.get(self._split, 0),
             n_classes=2,
             first_and_last_labels=("no_tumor", "tumor"),
         )

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/__init__.py

@@ -13,10 +13,11 @@ from eva.vision.data.transforms.intensity import (
     RandShiftIntensity,
     ScaleIntensityRange,
 )
-from eva.vision.data.transforms.spatial import RandFlip, RandRotate90, Spacing
+from eva.vision.data.transforms.spatial import RandFlip, RandRotate90, Resize, Spacing
 from eva.vision.data.transforms.utility import EnsureChannelFirst
 
 __all__ = [
+    "Resize",
     "ResizeAndCrop",
     "Squeeze",
     "CropForeground",

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"]