kaiko-eva 0.1.8__py3-none-any.whl → 0.2.1__py3-none-any.whl

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

@@ -0,0 +1,55 @@
+"""Intensity shifting transform."""
+
+import functools
+from typing import Any, Dict
+
+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
+
+
+class RandShiftIntensity(v2.Transform):
+    """Randomly shift intensity with randomly picked offset."""
+
+    def __init__(
+        self,
+        offsets: tuple[float, float] | float,
+        safe: bool = False,
+        prob: float = 0.1,
+        channel_wise: bool = False,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            offsets: Offset range to randomly shift.
+                if single number, offset value is picked from (-offsets, offsets).
+            safe: If `True`, then do safe dtype convert when intensity overflow.
+                E.g., `[256, -12]` -> `[array(0), array(244)]`. If `True`, then
+                `[256, -12]` -> `[array(255), array(0)]`.
+            prob: Probability of shift.
+            channel_wise: If True, shift intensity on each channel separately.
+                For each channel, a random offset will be chosen. Please ensure
+                that the first dimension represents the channel of the image if True.
+        """
+        super().__init__()
+
+        self._rand_swift_intensity = monai_intensity_transforms.RandShiftIntensity(
+            offsets=offsets,
+            safe=safe,
+            prob=prob,
+            channel_wise=channel_wise,
+        )
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_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_swift_intensity(inpt)
+        return tv_tensors.wrap(inpt_scaled, like=inpt)

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

@@ -0,0 +1,56 @@
+"""Intensity scaling transform."""
+
+import functools
+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
+
+
+class ScaleIntensityRange(v2.Transform):
+    """Intensity scaling transform.
+
+    Scaling from [a_min, a_max] to [b_min, b_max] with clip option.
+
+    When `b_min` or `b_max` are `None`, `scaled_array * (b_max - b_min) + b_min`
+    will be skipped. If `clip=True`, when `b_min`/`b_max` is None, the clipping
+    is not performed on the corresponding edge.
+    """
+
+    def __init__(
+        self,
+        input_range: Tuple[float, float],
+        output_range: Tuple[float, float] | None = None,
+        clip: bool = True,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            input_range: Intensity original range min and max.
+            output_range: Intensity target range min and max.
+            clip: Whether to perform clip after scaling.
+        """
+        super().__init__()
+
+        self._scale_intensity_range = monai_intensity_transforms.ScaleIntensityRange(
+            a_min=input_range[0],
+            a_max=input_range[1],
+            b_min=output_range[0] if output_range else None,
+            b_max=output_range[1] if output_range else None,
+            clip=clip,
+        )
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_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

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

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

@@ -0,0 +1,72 @@
+"""Flip transforms."""
+
+import functools
+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
+
+
+class RandFlip(v2.Transform):
+    """Randomly flips the image along axes."""
+
+    def __init__(
+        self,
+        prob: float = 0.1,
+        spatial_axes: Sequence[int] | int | None = None,
+        apply_per_axis: bool = True,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            prob: Probability of flipping.
+            spatial_axes: Spatial axes along which to flip over. Default is None.
+            apply_per_axis: If True, will apply a random flip transform to each
+                axis individually (if spatial_axes is a sequence of multiple axis).
+                If False, will apply a single random flip transform applied to all axes.
+        """
+        super().__init__()
+
+        if apply_per_axis:
+            if not isinstance(spatial_axes, (list, tuple)):
+                raise ValueError(
+                    "`spatial_axis` is expected to be sequence `apply_per_axis` "
+                    f"is enabled, got {type(spatial_axes)}"
+                )
+            self._flips = [
+                monai_spatial_transforms.RandFlip(prob=prob, spatial_axis=axis)
+                for axis in spatial_axes
+            ]
+        else:
+            self._flips = [monai_spatial_transforms.RandFlip(prob=prob, spatial_axis=spatial_axes)]
+
+    def _get_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:
+        return inpt
+
+    @_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)
+    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)
+
+    def _apply_flips(self, inpt: Any) -> Any:
+        for flip in self._flips:
+            inpt = flip(img=inpt, randomize=False)
+        return inpt

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

@@ -0,0 +1,53 @@
+"""Rotation transforms."""
+
+import functools
+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
+
+
+class RandRotate90(v2.Transform):
+    """Rotate input tensors by 90 degrees."""
+
+    def __init__(
+        self,
+        prob: float = 0.1,
+        max_k: int = 3,
+        spatial_axes: tuple[int, int] = (1, 2),
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            prob: probability of rotating.
+                (Default 0.1, with 10% probability it returns a rotated array)
+            max_k: number of rotations will be sampled from `np.random.randint(max_k) + 1`.
+            spatial_axes: 2 int numbers, defines the plane to rotate with 2 spatial axes.
+                Default: (1, 2), so for [C, T, H, W] will rotate along (H, W) plane (MONAI ignores
+                the first C dimension).
+        """
+        super().__init__()
+
+        self._rotate = monai_spatial_transforms.RandRotate90(
+            prob=prob, max_k=max_k, spatial_axes=spatial_axes
+        )
+
+    def _get_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:
+        return inpt
+
+    @_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

@@ -0,0 +1,69 @@
+"""Spacing resample transform."""
+
+import functools
+from typing import Any, Dict, List, Sequence
+
+import numpy as np
+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
+
+
+class Spacing(v2.Transform):
+    """Resample input image into the specified `pixdim`.
+
+    - Expects tensors of shape `[C, T, H, W]`.
+    """
+
+    def __init__(
+        self,
+        pixdim: Sequence[float] | float | np.ndarray,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            pixdim: output voxel spacing. if providing a single number,
+                will use it for the first dimension. Items of the pixdim
+                sequence map to the spatial dimensions of input image, if
+                length of pixdim sequence is longer than image spatial
+                dimensions, will ignore the longer part, if shorter, will
+                pad with the last value. For example, for 3D image if pixdim
+                is [1.0, 2.0] it will be padded to [1.0, 2.0, 2.0] if the
+                components of the `pixdim` are non-positive values, the
+                transform will use the corresponding components of the original
+                pixdim, which is computed from the `affine` matrix of input image.
+        """
+        super().__init__()
+
+        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]:
+        self._affine = next(
+            inpt.affine for inpt in flat_inputs if isinstance(inpt, eva_tv_tensors.Volume)
+        )
+        return {}
+
+    @functools.singledispatchmethod
+    @override
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return inpt
+
+    @_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)
+    def _(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        inpt_spacing = self._spacing(
+            meta_tensor.MetaTensor(inpt, affine=self._affine), mode="nearest"
+        )
+        return tv_tensors.wrap(inpt_spacing.to(dtype=torch.long), like=inpt)

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

@@ -0,0 +1,5 @@
+"""Transforms for utility operations."""
+
+from eva.vision.data.transforms.utility.ensure_channel_first import EnsureChannelFirst
+
+__all__ = ["EnsureChannelFirst"]

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

@@ -0,0 +1,51 @@
+"""Adjust or add the channel dimension of input data to ensure `channel_first` shape."""
+
+import functools
+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
+
+
+class EnsureChannelFirst(v2.Transform):
+    """Adjust or add the channel dimension of input data to ensure `channel_first` shape."""
+
+    def __init__(
+        self,
+        strict_check: bool = True,
+        channel_dim: None | str | int = None,
+    ) -> None:
+        """Initializes the transform.
+
+        Args:
+            strict_check: whether to raise an error when the meta information is insufficient.
+            channel_dim: This argument can be used to specify the original channel dimension
+                    (integer) of the input array.
+                It overrides the `original_channel_dim` from provided MetaTensor input.
+                If the input array doesn't have a channel dim, this value should be
+                    ``'no_channel'``.
+                If this is set to `None`, this class relies on `img` or `meta_dict` to provide
+                    the channel dimension.
+        """
+        super().__init__()
+
+        self._ensure_channel_first = monai_utility_transforms.EnsureChannelFirst(
+            strict_check=strict_check,
+            channel_dim=channel_dim,
+        )
+
+    @functools.singledispatchmethod
+    @override
+    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)
+    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/data/tv_tensors/__init__.py

@@ -0,0 +1,5 @@
+"""Custom `tv_tensors` types for torchvision."""
+
+from eva.vision.data.tv_tensors.volume import Volume
+
+__all__ = ["Volume"]

eva/vision/data/tv_tensors/volume.py

@@ -0,0 +1,61 @@
+"""Custom `tv_tensors` type for 3D Volumes."""
+
+from typing import Any, Dict, Optional, Union
+
+import torch
+from monai.data import meta_tensor
+from torchvision import tv_tensors
+from typing_extensions import override
+
+
+class Volume(tv_tensors.Video):
+    """:class:`torchvision.TVTensor` subclass for 3D volumes.
+
+    - Adds optional metadata and affine matrix to the tensor.
+    - Expects tensors to be of shape `[..., T, C, H, W]`.
+
+    Args:
+        data: Any data that can be turned into a tensor with :func:`torch.as_tensor`.
+        affine: Affine matrix of the volume. Expected to be of shape `[4, 4]`, and
+            columns to correspond to [T, H, W, (translation)] dimensions. Note that
+            `nibabel` by default uses [H, W, T, (translation)] order for affine matrices.
+        metadata: Metadata associated with the volume.
+        dtype: Desired data type. If omitted, will be inferred from `data`.
+        device: Desired device.
+        requires_grad: Whether autograd should record operations.
+    """
+
+    @override
+    def __new__(
+        cls,
+        data: Any,
+        affine: torch.Tensor | None = None,
+        metadata: Dict[str, Any] | None = None,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[Union[torch.device, str, int]] = None,
+        requires_grad: Optional[bool] = None,
+    ) -> "Volume":
+        cls.affine = affine
+        cls.metadata = metadata
+
+        return super().__new__(cls, data, dtype=dtype, device=device, requires_grad=requires_grad)  # type: ignore
+
+    @classmethod
+    def from_meta_tensor(cls, meta_tensor: meta_tensor.MetaTensor) -> "Volume":
+        """Creates an instance from a :class:`monai.data.meta_tensor.MetaTensor`."""
+        return cls(
+            meta_tensor.data,
+            affine=meta_tensor.affine,
+            metadata=meta_tensor.meta,
+            dtype=meta_tensor.dtype,
+            device=meta_tensor.device,
+            requires_grad=meta_tensor.requires_grad,
+        )  # type: ignore
+
+    def to_meta_tensor(self) -> meta_tensor.MetaTensor:
+        """Converts the volume to a :class:`monai.data.meta_tensor.MetaTensor`."""
+        return meta_tensor.MetaTensor(self, affine=self.affine, meta=self.metadata)
+
+    def __repr__(self, *, tensor_contents: Any = None) -> str:
+        """Returns the string representation of the object."""
+        return self._make_repr()

eva/vision/metrics/segmentation/monai_dice.py

@@ -1,5 +1,7 @@
 """Wrapper for dice score metric from MONAI."""
 
+from typing import Literal
+
 from monai.metrics.meandice import DiceMetric
 from typing_extensions import override
 
@@ -14,6 +16,7 @@ class MonaiDiceScore(wrappers.MonaiMetricWrapper):
         self,
         num_classes: int,
         include_background: bool = True,
+        input_format: Literal["one-hot", "index"] = "index",
         reduction: str = "mean",
         ignore_index: int | None = None,
         **kwargs,
@@ -24,6 +27,8 @@ class MonaiDiceScore(wrappers.MonaiMetricWrapper):
             num_classes: The number of classes in the dataset.
             include_background: Whether to include the background class in the computation.
             reduction: The method to reduce the dice score. Options are `"mean"`, `"sum"`, `"none"`.
+            input_format: Choose between "one-hot" for one-hot encoded tensors or "index"
+                for index tensors.
             ignore_index: Integer specifying a target class to ignore. If given, this class
                 index does not contribute to the returned score.
             kwargs: Additional keyword arguments for instantiating monai's `DiceMetric` class.
@@ -40,11 +45,13 @@ class MonaiDiceScore(wrappers.MonaiMetricWrapper):
         self.reduction = reduction
         self.orig_num_classes = num_classes
         self.ignore_index = ignore_index
+        self.input_format = input_format
 
     @override
     def update(self, preds, target):
-        preds = _utils.index_to_one_hot(preds, num_classes=self.orig_num_classes)
-        target = _utils.index_to_one_hot(target, num_classes=self.orig_num_classes)
+        if self.input_format == "index":
+            preds = _utils.index_to_one_hot(preds, num_classes=self.orig_num_classes)
+            target = _utils.index_to_one_hot(target, num_classes=self.orig_num_classes)
         if self.ignore_index is not None:
             preds, target = _utils.apply_ignore_index(preds, target, self.ignore_index)
         return super().update(preds, target)

eva/vision/models/modules/semantic_segmentation.py

@@ -1,10 +1,11 @@
 """"Neural Network Semantic Segmentation Module."""
 
-from typing import Any, Callable, Dict, Iterable, List, Tuple
+from typing import Any, Callable, Dict, Iterable, List
 
 import torch
 from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
 from lightning.pytorch.utilities.types import STEP_OUTPUT
+from monai.inferers.inferer import Inferer
 from torch import nn, optim
 from torch.optim import lr_scheduler
 from typing_extensions import override
@@ -15,6 +16,7 @@ from eva.core.models.modules.typings import INPUT_BATCH, INPUT_TENSOR_BATCH
 from eva.core.models.modules.utils import batch_postprocess, grad, submodule_state_dict
 from eva.core.utils import parser
 from eva.vision.models.networks import decoders
+from eva.vision.models.networks.decoders import segmentation
 from eva.vision.models.networks.decoders.segmentation.typings import DecoderInputs
 
 
@@ -23,10 +25,11 @@ class SemanticSegmentationModule(module.ModelModule):
 
     def __init__(
         self,
-        decoder: decoders.Decoder,
+        decoder: decoders.Decoder | nn.Module,
         criterion: Callable[..., torch.Tensor],
         encoder: Dict[str, Any] | Callable[[torch.Tensor], List[torch.Tensor]] | None = None,
         lr_multiplier_encoder: float = 0.0,
+        inferer: Inferer | None = None,
         optimizer: OptimizerCallable = optim.AdamW,
         lr_scheduler: LRSchedulerCallable = lr_scheduler.ConstantLR,
         metrics: metrics_lib.MetricsSchema | None = None,
@@ -44,6 +47,8 @@ class SemanticSegmentationModule(module.ModelModule):
                 during the `configure_model` step.
             lr_multiplier_encoder: The learning rate multiplier for the
                 encoder parameters. If `0`, it will freeze the encoder.
+            inferer: An optional MONAI `Inferer` for inference
+                postprocess during evaluation.
             optimizer: The optimizer to use.
             lr_scheduler: The learning rate scheduler to use.
             metrics: The metric groups to track.
@@ -62,6 +67,7 @@ class SemanticSegmentationModule(module.ModelModule):
         self.optimizer = optimizer
         self.lr_scheduler = lr_scheduler
         self.save_decoder_only = save_decoder_only
+        self.inferer = inferer
 
     @override
     def configure_model(self) -> None:
@@ -104,25 +110,15 @@ class SemanticSegmentationModule(module.ModelModule):
     @override
     def forward(
         self,
-        inputs: torch.Tensor,
-        to_size: Tuple[int, int] | None = None,
+        tensor: torch.Tensor,
         *args: Any,
         **kwargs: Any,
     ) -> torch.Tensor:
-        """Maps the input tensor (image tensor or embeddings) to masks.
-
-        If `inputs` is image tensor, then the `self.encoder`
-        should be implemented, otherwise it will be interpreted
-        as embeddings, where the `to_size` should be given.
-        """
-        if self.encoder is None and to_size is None:
-            raise ValueError(
-                "Please provide the expected `to_size` that the "
-                "decoder should map the embeddings (`inputs`) to."
-            )
-        features = self.encoder(inputs) if self.encoder else inputs
-        decoder_inputs = DecoderInputs(features, to_size or inputs.shape[-2:], inputs)  # type: ignore
-        return self.decoder(decoder_inputs)
+        return (
+            self.inferer(tensor, network=self._forward_networks)
+            if self.inferer is not None and not self.training
+            else self._forward_networks(tensor)
+        )
 
     @override
     def training_step(self, batch: INPUT_TENSOR_BATCH, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
@@ -137,7 +133,9 @@ class SemanticSegmentationModule(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 self.encoder(tensor) if isinstance(self.encoder, nn.Module) else tensor
 
@@ -170,7 +168,7 @@ class SemanticSegmentationModule(module.ModelModule):
             The batch step output.
         """
         data, targets, metadata = INPUT_TENSOR_BATCH(*batch)
-        predictions = self(data, to_size=targets.shape[-2:])
+        predictions = self(data)
         loss = self.criterion(predictions, targets)
         return {
             "loss": loss,
@@ -178,3 +176,13 @@ class SemanticSegmentationModule(module.ModelModule):
             "predictions": predictions,
             "metadata": metadata,
         }
+
+    def _forward_networks(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Passes the input tensor through the encoder and decoder."""
+        features = self.encoder(tensor) if self.encoder else tensor
+        if isinstance(self.decoder, segmentation.Decoder):
+            if not isinstance(features, list):
+                raise ValueError(f"Expected a list of feature map tensors, got {type(features)}.")
+            image_size = (tensor.shape[-2], tensor.shape[-1])
+            return self.decoder(DecoderInputs(features, image_size, tensor))
+        return self.decoder(features)

eva/vision/models/networks/backbones/__init__.py

@@ -1,6 +1,13 @@
 """Vision Model Backbones API."""
 
-from eva.vision.models.networks.backbones import pathology, timm, torchhub, universal
+from eva.vision.models.networks.backbones import pathology, radiology, timm, universal
 from eva.vision.models.networks.backbones.registry import BackboneModelRegistry, register_model
 
-__all__ = ["pathology", "timm", "torchhub", "universal", "BackboneModelRegistry", "register_model"]
+__all__ = [
+    "radiology",
+    "pathology",
+    "timm",
+    "universal",
+    "BackboneModelRegistry",
+    "register_model",
+]

eva/vision/models/networks/backbones/pathology/__init__.py

@@ -1,9 +1,14 @@
 """Vision Pathology Model Backbones API."""
 
-from eva.vision.models.networks.backbones.pathology.bioptimus import bioptimus_h_optimus_0
+from eva.vision.models.networks.backbones.pathology.bioptimus import (
+    bioptimus_h0_mini,
+    bioptimus_h_optimus_0,
+)
 from eva.vision.models.networks.backbones.pathology.gigapath import prov_gigapath
 from eva.vision.models.networks.backbones.pathology.histai import histai_hibou_b, histai_hibou_l
+from eva.vision.models.networks.backbones.pathology.hkust import hkust_gpfm
 from eva.vision.models.networks.backbones.pathology.kaiko import (
+    kaiko_midnight_12k,
     kaiko_vitb8,
     kaiko_vitb16,
     kaiko_vitl14,
@@ -11,11 +16,12 @@ from eva.vision.models.networks.backbones.pathology.kaiko import (
     kaiko_vits16,
 )
 from eva.vision.models.networks.backbones.pathology.lunit import lunit_vits8, lunit_vits16
-from eva.vision.models.networks.backbones.pathology.mahmood import mahmood_uni
+from eva.vision.models.networks.backbones.pathology.mahmood import mahmood_uni, mahmood_uni2_h
 from eva.vision.models.networks.backbones.pathology.owkin import owkin_phikon, owkin_phikon_v2
 from eva.vision.models.networks.backbones.pathology.paige import paige_virchow2
 
 __all__ = [
+    "kaiko_midnight_12k",
     "kaiko_vitb16",
     "kaiko_vitb8",
     "kaiko_vitl14",
@@ -26,9 +32,12 @@ __all__ = [
     "lunit_vits16",
     "lunit_vits8",
     "mahmood_uni",
+    "mahmood_uni2_h",
     "bioptimus_h_optimus_0",
+    "bioptimus_h0_mini",
     "prov_gigapath",
     "histai_hibou_b",
     "histai_hibou_l",
     "paige_virchow2",
+    "hkust_gpfm",
 ]