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

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

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

@@ -56,19 +56,20 @@ class RandCropByLabelClasses(base.RandomMonaiTransform):
     def set_random_state(self, seed: int) -> None:
         self._rand_crop.set_random_state(seed)
 
-    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+    @override
+    def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
         mask = next(inpt for inpt in flat_inputs if isinstance(inpt, tv_tensors.Mask))
         self._rand_crop.randomize(label=mask)
         return {}
 
     @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_crops = self._rand_crop(img=inpt, randomize=False)
         return [tv_tensors.wrap(crop, like=inpt) for crop in inpt_foreground_crops]

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

@@ -95,19 +95,20 @@ class RandCropByPosNegLabel(base.RandomMonaiTransform):
     def set_random_state(self, seed: int) -> None:
         self._rand_crop.set_random_state(seed)
 
-    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+    @override
+    def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
         mask = next(inpt for inpt in flat_inputs if isinstance(inpt, tv_tensors.Mask))
         self._rand_crop.randomize(label=mask)
         return {}
 
     @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_crops = self._rand_crop(img=inpt, randomize=False)
         return [tv_tensors.wrap(crop, like=inpt) for crop in inpt_foreground_crops]

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

@@ -5,14 +5,14 @@ from typing import Any, Dict, List, Sequence, Tuple
 
 from monai.transforms.croppad import array as monai_croppad_transforms
 from torchvision import tv_tensors
-from torchvision.transforms import v2
 from torchvision.transforms.v2 import _utils as tv_utils
 from typing_extensions import override
 
 from eva.vision.data import tv_tensors as eva_tv_tensors
+from eva.vision.data.transforms import base
 
 
-class RandSpatialCrop(v2.Transform):
+class RandSpatialCrop(base.TorchvisionTransformV2):
     """Crop image with random size or specific size ROI.
 
     It can crop at a random position as center or at the image center.
@@ -62,19 +62,20 @@ class RandSpatialCrop(v2.Transform):
         """Set the random state for the transform."""
         self._rand_spatial_crop.set_random_state(seed)
 
-    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+    @override
+    def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
         t, h, w = tv_utils.query_chw(flat_inputs)
         self._rand_spatial_crop.randomize((t, h, w))
         return {}
 
     @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:
         slices = self._get_crop_slices()
         inpt_rand_crop = self._cropper(inpt, slices=slices)

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

@@ -6,13 +6,13 @@ from typing import Any, Dict, Sequence
 from monai.transforms.croppad import array as monai_croppad_transforms
 from monai.utils.enums import Method, 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 SpatialPad(v2.Transform):
+class SpatialPad(base.TorchvisionTransformV2):
     """Performs padding to the data.
 
     Padding is applied symmetric for all sides or all on one side for each dimension.
@@ -56,12 +56,12 @@ class SpatialPad(v2.Transform):
 
     @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_padded = self._spatial_pad(inpt)
         return tv_tensors.wrap(inpt_padded, like=inpt)

eva/vision/data/transforms/intensity/rand_scale_intensity.py

@@ -53,11 +53,11 @@ class RandScaleIntensity(base.RandomMonaiTransform):
 
     @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.Image)
+    @transform.register(eva_tv_tensors.Volume)
     def _(self, inpt: tv_tensors.Image, params: Dict[str, Any]) -> Any:
         inpt_scaled = self._rand_scale_intensity(inpt)
         return tv_tensors.wrap(inpt_scaled, like=inpt)

eva/vision/data/transforms/intensity/rand_shift_intensity.py

@@ -49,11 +49,11 @@ class RandShiftIntensity(base.RandomMonaiTransform):
 
     @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.Image)
+    @transform.register(eva_tv_tensors.Volume)
     def _(self, inpt: tv_tensors.Image, params: Dict[str, Any]) -> Any:
         inpt_scaled = self._rand_shift_intensity(inpt)
         return tv_tensors.wrap(inpt_scaled, like=inpt)

eva/vision/data/transforms/intensity/scale_intensity_ranged.py

@@ -5,13 +5,13 @@ from typing import Any, Dict, Tuple
 
 from monai.transforms.intensity import array as monai_intensity_transforms
 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 ScaleIntensityRange(v2.Transform):
+class ScaleIntensityRange(base.TorchvisionTransformV2):
     """Intensity scaling transform.
 
     Scaling from [a_min, a_max] to [b_min, b_max] with clip option.
@@ -46,11 +46,11 @@ class ScaleIntensityRange(v2.Transform):
 
     @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.Image)
+    @transform.register(eva_tv_tensors.Volume)
     def _(self, inpt: tv_tensors.Image, params: Dict[str, Any]) -> Any:
         inpt_scaled = self._scale_intensity_range(inpt)
         return tv_tensors.wrap(inpt_scaled, like=inpt)

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

@@ -1,7 +1,8 @@
 """Transforms for spatial operations."""
 
 from eva.vision.data.transforms.spatial.flip import RandFlip
+from eva.vision.data.transforms.spatial.resize import Resize
 from eva.vision.data.transforms.spatial.rotate import RandRotate90
 from eva.vision.data.transforms.spatial.spacing import Spacing
 
-__all__ = ["Spacing", "RandFlip", "RandRotate90"]
+__all__ = ["Spacing", "RandFlip", "RandRotate90", "Resize"]

eva/vision/data/transforms/spatial/flip.py

@@ -6,13 +6,13 @@ from typing import Any, Dict, List, Sequence
 import torch
 from monai.transforms.spatial import array as monai_spatial_transforms
 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 RandFlip(v2.Transform):
+class RandFlip(base.TorchvisionTransformV2):
     """Randomly flips the image along axes."""
 
     def __init__(
@@ -45,23 +45,24 @@ class RandFlip(v2.Transform):
         else:
             self._flips = [monai_spatial_transforms.RandFlip(prob=prob, spatial_axis=spatial_axes)]
 
-    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+    @override
+    def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
         for flip in self._flips:
             flip.randomize(None)
         return {}
 
     @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.Image)
+    @transform.register(eva_tv_tensors.Volume)
     def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
         inpt_flipped = self._apply_flips(inpt)
         return tv_tensors.wrap(inpt_flipped, like=inpt)
 
-    @_transform.register(tv_tensors.Mask)
+    @transform.register(tv_tensors.Mask)
     def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
         inpt_flipped = torch.tensor(self._apply_flips(inpt), dtype=torch.long)
         return tv_tensors.wrap(inpt_flipped, like=inpt)

eva/vision/data/transforms/spatial/functional/__init__.py

@@ -0,0 +1,5 @@
+"""Functional API for spatial transforms."""
+
+from eva.vision.data.transforms.spatial.functional.resize import resize_to_max_bytes
+
+__all__ = ["resize_to_max_bytes"]

eva/vision/data/transforms/spatial/functional/resize.py

@@ -0,0 +1,26 @@
+"""Functional resizing utilities."""
+
+import io
+from typing import Tuple
+
+from PIL import Image
+from torchvision import tv_tensors
+from torchvision.transforms.v2 import functional as F
+
+
+def resize_to_max_bytes(image: tv_tensors.Image, max_bytes: int) -> tv_tensors.Image:
+    """Resize the image to fit within the specified byte size."""
+    image_pil = F.to_pil_image(image)
+    image_bytes = io.BytesIO()
+    image_pil.save(image_bytes, format="PNG", optimize=True)
+
+    while image_bytes.tell() > max_bytes:
+        size: Tuple[int, int] = image_pil.size  # type: ignore
+        w, h = size
+        scale = (max_bytes / image_bytes.tell()) ** 0.5
+        new_size = (max(1, int(h * scale)), max(1, int(w * scale)))
+        image_pil = image_pil.resize(new_size, Image.Resampling.LANCZOS)
+        image_bytes = io.BytesIO()
+        image_pil.save(image_bytes, format="PNG", optimize=True)
+
+    return tv_tensors.Image(F.pil_to_tensor(image_pil))

eva/vision/data/transforms/spatial/resize.py

@@ -0,0 +1,63 @@
+"""Image resize transforms."""
+
+import functools
+from typing import Any, Dict
+
+from torchvision import tv_tensors
+from torchvision.transforms import v2
+from typing_extensions import override
+
+from eva.vision.data.transforms import base
+from eva.vision.data.transforms.spatial import functional
+
+
+class Resize(base.TorchvisionTransformV2):
+    """Resize transform for images with spatial or byte-based constraints.
+
+    This transform provides two mutually exclusive modes of resizing:
+    1. Spatial resizing: Resize to a specific (height, width) dimension
+    2. Byte-based resizing: Resize to fit within a maximum byte size
+
+    The latter is particularly useful for API models (e.g. Claude 3.7) that
+    have strict byte size limits for image inputs.
+    """
+
+    def __init__(self, size: tuple[int, int] | None = None, max_bytes: int | None = None) -> None:
+        """Initializes the transform.
+
+        Args:
+            size: Target size as (height, width) tuple for spatial resizing.
+                If provided, max_bytes must be None.
+            max_bytes: Maximum allowed byte size for the image.
+                If provided, size must be None. Must be a positive integer.
+
+        Raises:
+            ValueError: If both size and max_bytes are provided, or if max_bytes
+                is not a positive integer.
+        """
+        if size is not None and max_bytes is not None:
+            raise ValueError("Cannot provide both 'size' and 'max_bytes' parameters.")
+        if max_bytes is not None and max_bytes <= 0:
+            raise ValueError("'max_bytes' must be a positive integer.")
+
+        super().__init__()
+
+        self.size = size
+        self.max_bytes = max_bytes
+        self.resize_fn = None
+
+        if size is not None:
+            self.resize_fn = v2.Resize(size=size)
+        elif max_bytes is not None:
+            self.resize_fn = functools.partial(functional.resize_to_max_bytes, max_bytes=max_bytes)
+
+    @functools.singledispatchmethod
+    @override
+    def transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @transform.register(tv_tensors.Image)
+    @transform.register(tv_tensors.Mask)
+    def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        inpt_resized = self.resize_fn(inpt) if self.resize_fn is not None else inpt
+        return tv_tensors.wrap(inpt_resized, like=inpt)

eva/vision/data/transforms/spatial/rotate.py

@@ -5,13 +5,13 @@ from typing import Any, Dict, List
 
 from monai.transforms.spatial import array as monai_spatial_transforms
 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 RandRotate90(v2.Transform):
+class RandRotate90(base.TorchvisionTransformV2):
     """Rotate input tensors by 90 degrees."""
 
     def __init__(
@@ -36,18 +36,19 @@ class RandRotate90(v2.Transform):
             prob=prob, max_k=max_k, spatial_axes=spatial_axes
         )
 
-    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+    @override
+    def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
         self._rotate.randomize()
         return {}
 
     @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_rotated = self._rotate(img=inpt, randomize=False)
         return tv_tensors.wrap(inpt_rotated, like=inpt)

eva/vision/data/transforms/spatial/spacing.py

@@ -8,13 +8,13 @@ import torch
 from monai.data import meta_tensor
 from monai.transforms.spatial import array as monai_spatial_transforms
 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 Spacing(v2.Transform):
+class Spacing(base.TorchvisionTransformV2):
     """Resample input image into the specified `pixdim`.
 
     - Expects tensors of shape `[C, T, H, W]`.
@@ -43,7 +43,8 @@ class Spacing(v2.Transform):
         self._spacing = monai_spatial_transforms.Spacing(pixdim=pixdim, recompute_affine=True)
         self._affine = None
 
-    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+    @override
+    def make_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
         self._affine = next(
             inpt.affine for inpt in flat_inputs if isinstance(inpt, eva_tv_tensors.Volume)
         )
@@ -51,17 +52,17 @@ class Spacing(v2.Transform):
 
     @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(eva_tv_tensors.Volume)
+    @transform.register(eva_tv_tensors.Volume)
     def _(self, inpt: eva_tv_tensors.Volume, params: Dict[str, Any]) -> Any:
         inpt_spacing = self._spacing(inpt.to_meta_tensor(), mode="bilinear")
         if not isinstance(inpt_spacing, meta_tensor.MetaTensor):
             raise ValueError(f"Expected MetaTensor, got {type(inpt_spacing)}")
         return eva_tv_tensors.Volume.from_meta_tensor(inpt_spacing)
 
-    @_transform.register(tv_tensors.Mask)
+    @transform.register(tv_tensors.Mask)
     def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
         inpt_spacing = self._spacing(
             meta_tensor.MetaTensor(inpt, affine=self._affine), mode="nearest"

eva/vision/data/transforms/utility/ensure_channel_first.py

@@ -5,13 +5,13 @@ from typing import Any, Dict
 
 from monai.transforms.utility import array as monai_utility_transforms
 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 EnsureChannelFirst(v2.Transform):
+class EnsureChannelFirst(base.TorchvisionTransformV2):
     """Adjust or add the channel dimension of input data to ensure `channel_first` shape."""
 
     def __init__(
@@ -40,12 +40,12 @@ class EnsureChannelFirst(v2.Transform):
 
     @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_channel_first = self._ensure_channel_first(inpt)
         return tv_tensors.wrap(inpt_channel_first, like=inpt)

eva/vision/models/networks/backbones/universal/vit.py

@@ -54,6 +54,30 @@ def vit_small_patch16_224_dino(
     )
 
 
+@backbone_registry.register("universal/vit_tiny_patch16_224_random")
+def vit_tiny_patch16_224_random(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes a ViT-Tiny16 baseline model with random weights.
+
+    Args:
+        dynamic_img_size: Support different input image sizes by allowing to change
+            the grid size (interpolate abs and/or ROPE pos) in the forward pass.
+        out_indices: Whether and which multi-level patch embeddings to return.
+
+    Returns:
+        The torch ViTS-16 based foundation model.
+    """
+    return timm.create_model(
+        model_name="vit_tiny_patch16_224",
+        pretrained=False,
+        num_classes=0,
+        features_only=out_indices is not None,
+        out_indices=out_indices,
+        dynamic_img_size=dynamic_img_size,
+    )
+
+
 @backbone_registry.register("universal/vit_small_patch16_224_dino_1chan")
 def vit_small_patch16_224_dino_1chan(
     dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None

eva/vision/models/wrappers/from_registry.py

@@ -3,6 +3,7 @@
 from typing import Any, Callable, Dict
 
 import torch
+from torch import nn
 from typing_extensions import override
 
 from eva.core.models.wrappers import base
@@ -40,14 +41,14 @@ class ModelFromRegistry(base.BaseModel[torch.Tensor, torch.Tensor]):
         self._model_kwargs = model_kwargs or {}
         self._model_extra_kwargs = model_extra_kwargs or {}
 
-        self.load_model()
+        self.model = self.load_model()
 
     @override
-    def load_model(self) -> None:
-        self._model = factory.ModuleFactory(
+    def load_model(self) -> nn.Module:
+        ModelFromRegistry.__name__ = self._model_name
+
+        return factory.ModuleFactory(
             registry=backbone_registry,
             name=self._model_name,
             init_args=self._model_kwargs | self._model_extra_kwargs,
         )
-
-        ModelFromRegistry.__name__ = self._model_name