kaiko-eva 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

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,12 @@
 """"Neural Network Semantic Segmentation Module."""
 
+import functools
 from typing import Any, Callable, Dict, Iterable, List, Tuple
 
 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 +17,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,15 +26,17 @@ 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,
         postprocess: batch_postprocess.BatchPostProcess | None = None,
         save_decoder_only: bool = True,
+        spatial_dims: int = 2,
     ) -> None:
         """Initializes the neural net head module.
 
@@ -44,6 +49,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.
@@ -52,6 +59,8 @@ class SemanticSegmentationModule(module.ModelModule):
                 predictions and targets.
             save_decoder_only: Whether to save only the decoder during checkpointing. If False,
                 will also save the encoder (not recommended when frozen).
+            spatial_dims: The number of spatial dimensions, 2 for 2D
+                and 3 for 3D segmentation.
         """
         super().__init__(metrics=metrics, postprocess=postprocess)
 
@@ -62,6 +71,8 @@ class SemanticSegmentationModule(module.ModelModule):
         self.optimizer = optimizer
         self.lr_scheduler = lr_scheduler
         self.save_decoder_only = save_decoder_only
+        self.inferer = inferer
+        self.spatial_dims = spatial_dims
 
     @override
     def configure_model(self) -> None:
@@ -104,25 +115,16 @@ class SemanticSegmentationModule(module.ModelModule):
     @override
     def forward(
         self,
-        inputs: torch.Tensor,
-        to_size: Tuple[int, int] | None = None,
+        tensor: torch.Tensor,
+        to_size: Tuple[int, int],
         *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=functools.partial(self._forward_networks, to_size=to_size))
+            if self.inferer is not None and not self.training
+            else self._forward_networks(tensor, to_size=to_size)
+        )
 
     @override
     def training_step(self, batch: INPUT_TENSOR_BATCH, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
@@ -137,7 +139,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 +174,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, to_size=targets.shape[-self.spatial_dims :])
         loss = self.criterion(predictions, targets)
         return {
             "loss": loss,
@@ -178,3 +182,12 @@ class SemanticSegmentationModule(module.ModelModule):
             "predictions": predictions,
             "metadata": metadata,
         }
+
+    def _forward_networks(self, tensor: torch.Tensor, to_size: Tuple[int, int]) -> 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)}.")
+            return self.decoder(DecoderInputs(features, to_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",
 ]

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

@@ -5,6 +5,9 @@ from typing import Tuple
 import timm
 from torch import nn
 
+from eva.core.models import transforms
+from eva.vision.models import wrappers
+from eva.vision.models.networks.backbones import _utils
 from eva.vision.models.networks.backbones.registry import register_model
 
 
@@ -13,7 +16,9 @@ def bioptimus_h_optimus_0(
     dynamic_img_size: bool = True,
     out_indices: int | Tuple[int, ...] | None = None,
 ) -> nn.Module:
-    """Initializes the h_optimus_0 pathology FM by Bioptimus.
+    """Initializes the H-Optimus-0 pathology FM by Bioptimus.
+
+    See https://huggingface.co/bioptimus/H-optimus-0 for details.
 
     Args:
         dynamic_img_size: Whether to allow the interpolation embedding
@@ -32,3 +37,44 @@ def bioptimus_h_optimus_0(
         out_indices=out_indices,
         features_only=out_indices is not None,
     )
+
+
+@register_model("pathology/bioptimus_h0_mini")
+def bioptimus_h0_mini(
+    dynamic_img_size: bool = True,
+    out_indices: int | Tuple[int, ...] | None = None,
+    hf_token: str | None = None,
+    include_patch_tokens: bool = False,
+) -> nn.Module:
+    """Initializes H0-mini (ViT-B) pathology FM by Bioptimus.
+
+    This model was distilled from H-Optimus-0 on 40M TCGA tiles.
+
+    See https://huggingface.co/bioptimus/H0-mini for details.
+
+    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.
+        hf_token: HuggingFace token to download the model.
+        include_patch_tokens: Whether to combine the mean aggregated patch tokens with cls token.
+
+    Returns:
+        The model instance.
+    """
+    _utils.huggingface_login(hf_token)
+    return wrappers.TimmModel(
+        model_name="hf-hub:bioptimus/H0-mini",
+        out_indices=out_indices,
+        pretrained=True,
+        model_kwargs={
+            "dynamic_img_size": dynamic_img_size,
+            "mlp_layer": timm.layers.SwiGLUPacked,
+            "act_layer": nn.SiLU,
+        },
+        tensor_transforms=(
+            transforms.ExtractCLSFeatures(include_patch_tokens=include_patch_tokens)
+            if out_indices is None
+            else None
+        ),
+    )

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

@@ -0,0 +1,69 @@
+"""Pathology FMs from Hong Kong University of Science and Technology."""
+
+import re
+from typing import Tuple
+
+import timm
+from torch import nn
+
+from eva.core.models.wrappers import _utils
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+@register_model("pathology/hkust_gpfm")
+def hkust_gpfm(
+    dynamic_img_size: bool = True,
+    out_indices: int | Tuple[int, ...] | None = None,
+) -> nn.Module:
+    """Initializes GPFM model from Hong Kong University of Science and Technology.
+
+    Ma, J., Guo, Z., Zhou, F., Wang, Y., Xu, Y., et al. (2024).
+    Towards a generalizable pathology foundation model via unified knowledge
+    distillation (arXiv No. 2407.18449). arXiv. https://arxiv.org/abs/2407.18449
+
+    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 model instance.
+    """
+    return timm.create_model(
+        model_name="vit_large_patch14_dinov2",
+        pretrained=True,
+        pretrained_cfg={
+            "state_dict": _load_state_dict(),
+            "num_classes": 0,
+        },
+        out_indices=out_indices,
+        features_only=out_indices is not None,
+        **{
+            "img_size": 224,
+            "patch_size": 14,
+            "init_values": 1e-5,
+            "qkv_bias": True,
+            "dynamic_img_size": dynamic_img_size,
+        },
+    )
+
+
+def _load_state_dict() -> dict:
+    """Loads the state dict with model weights from github."""
+    state_dict = _utils.load_state_dict_from_url(
+        url="https://github.com/birkhoffkiki/GPFM/releases/download/ckpt/GPFM.pth",
+        md5="0dc7e345de84f385d09c8c782b4b3236",
+    )
+    return _convert_state_dict(state_dict["teacher"])
+
+
+def _convert_state_dict(state_dict: dict) -> dict:
+    """Rename state dict keys to match timm's format."""
+    state_dict = {
+        re.sub(r"blocks\.\d+\.(\d+)", r"blocks.\1", key.replace("backbone.", "")): value
+        for key, value in state_dict.items()
+    }
+    remove_keys = ["mask_token"] + [key for key in state_dict.keys() if "dino_head" in key]
+    for key in remove_keys:
+        state_dict.pop(key)
+    return state_dict

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

@@ -5,9 +5,27 @@ from typing import Tuple
 import torch
 from torch import nn
 
+from eva.vision.models.networks.backbones import _utils
 from eva.vision.models.networks.backbones.registry import register_model
 
 
+@register_model("pathology/kaiko_midnight_12k")
+def kaiko_midnight_12k(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
+    """Initializes the Midnight-12k pathology FM by kaiko.ai.
+
+    Args:
+        out_indices: Whether and which multi-level patch embeddings to return.
+
+    Returns:
+        The model instance.
+    """
+    return _utils.load_hugingface_model(
+        model_name="kaiko-ai/midnight",
+        out_indices=out_indices,
+        model_kwargs={"trust_remote_code": True},
+    )
+
+
 @register_model("pathology/kaiko_vits16")
 def kaiko_vits16(
     dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None

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

@@ -0,0 +1,11 @@
+"""Vision Radiology Model Backbones API."""
+
+from eva.vision.models.networks.backbones.radiology.swin_unetr import SwinUNETREncoder
+from eva.vision.models.networks.backbones.radiology.voco import VoCoB, VoCoH, VoCoL
+
+__all__ = [
+    "VoCoB",
+    "VoCoL",
+    "VoCoH",
+    "SwinUNETREncoder",
+]