kaiko-eva 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,62 @@
+"""Pathology FMs from MahmoodLab."""
+
+import os
+from pathlib import Path
+from typing import Tuple
+
+import huggingface_hub
+from loguru import logger
+from torch import nn
+
+from eva.vision.models import wrappers
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+@register_model("pathology/mahmood_uni")
+def mahmood_uni(
+    dynamic_img_size: bool = True,
+    out_indices: int | Tuple[int, ...] | None = None,
+    hf_token: str | None = None,
+    download_dir: str = os.path.join(str(Path.home()), ".cache/eva"),
+) -> nn.Module:
+    """Initializes UNI model from MahmoodLab.
+
+    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.
+        download_dir: Directory to download the model checkpoint.
+
+    Returns:
+        The model instance.
+    """
+    token = hf_token or os.environ.get("HF_TOKEN")
+    if not token:
+        raise ValueError(
+            "Please provide a HuggingFace token to download the model. "
+            "You can either pass it as an argument or set the env variable HF_TOKEN."
+        )
+
+    checkpoint_path = os.path.join(download_dir, "pytorch_model.bin")
+
+    if not os.path.exists(checkpoint_path):
+        logger.info(f"Downloading the model checkpoint to {download_dir} ...")
+        os.makedirs(download_dir, exist_ok=True)
+        huggingface_hub.login(token=token)
+        huggingface_hub.hf_hub_download(
+            "MahmoodLab/UNI",
+            filename="pytorch_model.bin",
+            local_dir=download_dir,
+            force_download=True,
+        )
+
+    return wrappers.TimmModel(
+        model_name="vit_large_patch16_224",
+        out_indices=out_indices,
+        model_kwargs={
+            "init_values": 1e-5,
+            "dynamic_img_size": dynamic_img_size,
+        },
+        checkpoint_path=checkpoint_path,
+    )

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

@@ -0,0 +1,22 @@
+"""Pathology FMs from owkin."""
+
+from typing import Tuple
+
+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/owkin_phikon")
+def owkin_phikon(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
+    """Initializes the phikon pathology FM by owkin (https://huggingface.co/owkin/phikon).
+
+    Args:
+        out_indices: Whether and which multi-level patch embeddings to return.
+            Currently only out_indices=1 is supported.
+
+    Returns:
+        The model instance.
+    """
+    return _utils.load_hugingface_model(model_name="owkin/phikon", out_indices=out_indices)

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

@@ -0,0 +1,47 @@
+"""Backbone Model Registry."""
+
+from typing import Any, Callable, Dict, List
+
+import torch.nn as nn
+
+
+class BackboneModelRegistry:
+    """A model registry for accessing backbone models by name."""
+
+    _registry: Dict[str, Callable[..., nn.Module]] = {}
+
+    @classmethod
+    def register(cls, name: str) -> Callable:
+        """Decorator to register a new model."""
+
+        def decorator(model_fn: Callable[..., nn.Module]) -> Callable[..., nn.Module]:
+            if name in cls._registry:
+                raise ValueError(f"Model {name} is already registered.")
+            cls._registry[name] = model_fn
+            return model_fn
+
+        return decorator
+
+    @classmethod
+    def get(cls, model_name: str) -> Callable[..., nn.Module]:
+        """Gets a model function from the registry."""
+        if model_name not in cls._registry:
+            raise ValueError(f"Model {model_name} not found in the registry.")
+        return cls._registry[model_name]
+
+    @classmethod
+    def load_model(cls, model_name: str, model_kwargs: Dict[str, Any] | None = None) -> nn.Module:
+        """Loads & initializes a model class from the registry."""
+        model_fn = cls.get(model_name)
+        return model_fn(**(model_kwargs or {}))
+
+    @classmethod
+    def list_models(cls) -> List[str]:
+        """List all models in the registry."""
+        register_models = [name for name in cls._registry.keys() if not name.startswith("timm")]
+        return register_models + ["timm/<model_name>"]
+
+
+def register_model(name: str) -> Callable:
+    """Simple decorator to register a model."""
+    return BackboneModelRegistry.register(name)

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

@@ -0,0 +1,5 @@
+"""timm backbones API."""
+
+from eva.vision.models.networks.backbones.timm.backbones import timm_model
+
+__all__ = ["timm_model"]

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

@@ -0,0 +1,54 @@
+"""timm backbones."""
+
+import functools
+from typing import Tuple
+
+import timm
+from loguru import logger
+from torch import nn
+
+from eva.vision.models import wrappers
+from eva.vision.models.networks.backbones.registry import BackboneModelRegistry
+
+
+def timm_model(
+    model_name: str,
+    checkpoint_path: str | None = None,
+    pretrained: bool = False,
+    dynamic_img_size: bool = True,
+    out_indices: int | Tuple[int, ...] | None = None,
+    **kwargs,
+) -> nn.Module:
+    """Initializes any ViT model from timm with weights from a specified checkpoint.
+
+    Args:
+        model_name: The name of the model to load.
+        checkpoint_path: The path to the checkpoint file.
+        pretrained: If set to `True`, load pretrained ImageNet-1k weights.
+        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.
+        **kwargs: Additional arguments to pass to the model
+
+    Returns:
+        The VIT model instance.
+    """
+    logger.info(
+        f"Loading timm model {model_name}"
+        + (f"using checkpoint {checkpoint_path}" if checkpoint_path else "")
+    )
+    return wrappers.TimmModel(
+        model_name=model_name,
+        checkpoint_path=checkpoint_path or "",
+        pretrained=pretrained,
+        out_indices=out_indices,
+        model_kwargs=kwargs,
+    )
+
+
+BackboneModelRegistry._registry.update(
+    {
+        f"timm/{model_name}": functools.partial(timm_model, model_name=model_name)
+        for model_name in timm.list_models()
+    }
+)

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

@@ -0,0 +1,8 @@
+"""Universal Vision Model Backbones API."""
+
+from eva.vision.models.networks.backbones.universal.vit import (
+    vit_small_patch16_224_dino,
+    vit_small_patch16_224_random,
+)
+
+__all__ = ["vit_small_patch16_224_dino", "vit_small_patch16_224_random"]

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

@@ -0,0 +1,54 @@
+"""Vision Transformers base universal backbones."""
+
+from typing import Tuple
+
+import timm
+from torch import nn
+
+from eva.vision.models.networks.backbones.registry import register_model
+
+
+@register_model("universal/vit_small_patch16_224_random")
+def vit_small_patch16_224_random(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes a ViTS-16 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_small_patch16_224.dino",
+        pretrained=False,
+        features_only=out_indices is not None,
+        out_indices=out_indices,
+        dynamic_img_size=dynamic_img_size,
+    )
+
+
+@register_model("universal/vit_small_patch16_224_dino")
+def vit_small_patch16_224_dino(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes a ViTS-16 baseline model pretrained w/ DINO.
+
+    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_small_patch16_224.dino",
+        pretrained=True,
+        features_only=out_indices is not None,
+        out_indices=out_indices,
+        dynamic_img_size=dynamic_img_size,
+    )

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

@@ -0,0 +1,6 @@
+"""Decoder heads API."""
+
+from eva.vision.models.networks.decoders import segmentation
+from eva.vision.models.networks.decoders.decoder import Decoder
+
+__all__ = ["segmentation", "Decoder"]

eva/vision/models/networks/decoders/decoder.py

@@ -0,0 +1,7 @@
+"""Semantic segmentation decoder base class."""
+
+from torch import nn
+
+
+class Decoder(nn.Module):
+    """Semantic segmentation decoder base class."""

eva/vision/models/networks/decoders/segmentation/__init__.py

@@ -0,0 +1,11 @@
+"""Segmentation decoder heads API."""
+
+from eva.vision.models.networks.decoders.segmentation.common import (
+    ConvDecoder1x1,
+    ConvDecoderMS,
+    SingleLinearDecoder,
+)
+from eva.vision.models.networks.decoders.segmentation.conv2d import ConvDecoder
+from eva.vision.models.networks.decoders.segmentation.linear import LinearDecoder
+
+__all__ = ["ConvDecoder1x1", "ConvDecoderMS", "SingleLinearDecoder", "ConvDecoder", "LinearDecoder"]

eva/vision/models/networks/decoders/segmentation/common.py

@@ -0,0 +1,74 @@
+"""Common semantic segmentation decoders.
+
+This module contains implementations of different types of decoder models 
+used in semantic segmentation. These decoders convert the high-level features 
+output by an encoder into pixel-wise predictions for segmentation tasks.
+"""
+
+from torch import nn
+
+from eva.vision.models.networks.decoders.segmentation import conv2d, linear
+
+
+class ConvDecoder1x1(conv2d.ConvDecoder):
+    """A convolutional decoder with a single 1x1 convolutional layer."""
+
+    def __init__(self, in_features: int, num_classes: int) -> None:
+        """Initializes the decoder.
+
+        Args:
+            in_features: The hidden dimension size of the embeddings.
+            num_classes: Number of output classes as channels.
+        """
+        super().__init__(
+            layers=nn.Conv2d(
+                in_channels=in_features,
+                out_channels=num_classes,
+                kernel_size=(1, 1),
+            ),
+        )
+
+
+class ConvDecoderMS(conv2d.ConvDecoder):
+    """A multi-stage convolutional decoder with upsampling and convolutional layers.
+
+    This decoder applies a series of upsampling and convolutional layers to transform
+    the input features into output predictions with the desired spatial resolution.
+
+    This decoder is based on the `+ms` segmentation decoder from DINOv2
+    (https://arxiv.org/pdf/2304.07193)
+    """
+
+    def __init__(self, in_features: int, num_classes: int) -> None:
+        """Initializes the decoder.
+
+        Args:
+            in_features: The hidden dimension size of the embeddings.
+            num_classes: Number of output classes as channels.
+        """
+        super().__init__(
+            layers=nn.Sequential(
+                nn.Upsample(scale_factor=2),
+                nn.Conv2d(in_features, 64, kernel_size=(3, 3), padding=(1, 1)),
+                nn.Upsample(scale_factor=2),
+                nn.Conv2d(64, num_classes, kernel_size=(3, 3), padding=(1, 1)),
+            ),
+        )
+
+
+class SingleLinearDecoder(linear.LinearDecoder):
+    """A simple linear decoder with a single fully connected layer."""
+
+    def __init__(self, in_features: int, num_classes: int) -> None:
+        """Initializes the decoder.
+
+        Args:
+            in_features: The hidden dimension size of the embeddings.
+            num_classes: Number of output classes as channels.
+        """
+        super().__init__(
+            layers=nn.Linear(
+                in_features=in_features,
+                out_features=num_classes,
+            ),
+        )

eva/vision/models/networks/decoders/segmentation/conv2d.py

@@ -0,0 +1,114 @@
+"""Convolutional based semantic segmentation decoder."""
+
+from typing import List, Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional
+
+from eva.vision.models.networks.decoders import decoder
+
+
+class ConvDecoder(decoder.Decoder):
+    """Convolutional segmentation decoder."""
+
+    def __init__(self, layers: nn.Module) -> None:
+        """Initializes the convolutional based decoder head.
+
+        Here the input nn layers will be directly applied to the
+        features of shape (batch_size, hidden_size, n_patches_height,
+        n_patches_width), where n_patches is image_size / patch_size.
+        Note the n_patches is also known as grid_size.
+
+        Args:
+            layers: The convolutional layers to be used as the decoder head.
+        """
+        super().__init__()
+
+        self._layers = layers
+
+    def _forward_features(self, features: List[torch.Tensor]) -> torch.Tensor:
+        """Forward function for multi-level feature maps to a single one.
+
+        It will interpolate the features and concat them into a single tensor
+        on the dimension axis of the hidden size.
+
+        Example:
+            >>> features = [torch.Tensor(16, 384, 14, 14), torch.Size(16, 384, 14, 14)]
+            >>> output = self._forward_features(features)
+            >>> assert output.shape == torch.Size([16, 768, 14, 14])
+
+        Args:
+            features: List of multi-level image features of shape (batch_size,
+                hidden_size, n_patches_height, n_patches_width).
+
+        Returns:
+            A tensor of shape (batch_size, hidden_size, n_patches_height,
+            n_patches_width) which is feature map of the decoder head.
+        """
+        if not isinstance(features, list) or features[0].ndim != 4:
+            raise ValueError(
+                "Input features should be a list of four (4) dimensional inputs of "
+                "shape (batch_size, hidden_size, n_patches_height, n_patches_width)."
+            )
+
+        upsampled_features = [
+            functional.interpolate(
+                input=embeddings,
+                size=features[0].shape[2:],
+                mode="bilinear",
+                align_corners=False,
+            )
+            for embeddings in features
+        ]
+        return torch.cat(upsampled_features, dim=1)
+
+    def _forward_head(self, patch_embeddings: torch.Tensor) -> torch.Tensor:
+        """Forward of the decoder head.
+
+        Args:
+            patch_embeddings: The patch embeddings tensor of shape
+                (batch_size, hidden_size, n_patches_height, n_patches_width).
+
+        Returns:
+            The logits as a tensor (batch_size, n_classes, upscale_height, upscale_width).
+        """
+        return self._layers(patch_embeddings)
+
+    def _cls_seg(
+        self,
+        logits: torch.Tensor,
+        image_size: Tuple[int, int],
+    ) -> torch.Tensor:
+        """Classify each pixel of the image.
+
+        Args:
+            logits: The decoder outputs of shape (batch_size, n_classes,
+                height, width).
+            image_size: The target image size (height, width).
+
+        Returns:
+            Tensor containing scores for all of the classes with shape
+            (batch_size, n_classes, image_height, image_width).
+        """
+        return functional.interpolate(logits, image_size, mode="bilinear")
+
+    def forward(
+        self,
+        features: List[torch.Tensor],
+        image_size: Tuple[int, int],
+    ) -> torch.Tensor:
+        """Maps the patch embeddings to a segmentation mask of the image size.
+
+        Args:
+            features: List of multi-level image features of shape (batch_size,
+                hidden_size, n_patches_height, n_patches_width).
+            image_size: The target image size (height, width).
+
+        Returns:
+            Tensor containing scores for all of the classes with shape
+            (batch_size, n_classes, image_height, image_width).
+        """
+        patch_embeddings = self._forward_features(features)
+        logits = self._forward_head(patch_embeddings)
+        return self._cls_seg(logits, image_size)

eva/vision/models/networks/decoders/segmentation/linear.py

@@ -0,0 +1,125 @@
+"""Linear based decoder."""
+
+from typing import List, Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional
+
+from eva.vision.models.networks.decoders import decoder
+
+
+class LinearDecoder(decoder.Decoder):
+    """Linear decoder."""
+
+    def __init__(self, layers: nn.Module) -> None:
+        """Initializes the linear based decoder head.
+
+        Here the input nn layers will be applied to the reshaped
+        features (batch_size, patch_embeddings, hidden_size) from
+        the input (batch_size, hidden_size, height, width) and then
+        unwrapped again to (batch_size, n_classes, height, width).
+
+        Args:
+            layers: The linear layers to be used as the decoder head.
+        """
+        super().__init__()
+
+        self._layers = layers
+
+    def _forward_features(self, features: List[torch.Tensor]) -> torch.Tensor:
+        """Forward function for multi-level feature maps to a single one.
+
+        It will interpolate the features and concat them into a single tensor
+        on the dimension axis of the hidden size.
+
+        Example:
+            >>> features = [torch.Tensor(16, 384, 14, 14), torch.Size(16, 384, 14, 14)]
+            >>> output = self._forward_features(features)
+            >>> assert output.shape == torch.Size([16, 768, 14, 14])
+
+        Args:
+            features: List of multi-level image features of shape (batch_size,
+                hidden_size, n_patches_height, n_patches_width).
+
+        Returns:
+            A tensor of shape (batch_size, hidden_size, n_patches_height,
+            n_patches_width) which is feature map of the decoder head.
+        """
+        if not isinstance(features, list) or features[0].ndim != 4:
+            raise ValueError(
+                "Input features should be a list of four (4) dimensional inputs of "
+                "shape (batch_size, hidden_size, n_patches_height, n_patches_width)."
+            )
+
+        upsampled_features = [
+            functional.interpolate(
+                input=embeddings,
+                size=features[0].shape[2:],
+                mode="bilinear",
+                align_corners=False,
+            )
+            for embeddings in features
+        ]
+        return torch.cat(upsampled_features, dim=1)
+
+    def _forward_head(self, patch_embeddings: torch.Tensor) -> torch.Tensor:
+        """Forward of the decoder head.
+
+        Here the following transformations will take place:
+        - (batch_size, hidden_size, n_patches_height, n_patches_width)
+        - (batch_size, hidden_size, n_patches_height * n_patches_width)
+        - (batch_size, n_patches_height * n_patches_width, hidden_size)
+        - (batch_size, n_patches_height * n_patches_width, n_classes)
+        - (batch_size, n_classes, n_patches_height, n_patches_width)
+
+        Args:
+            patch_embeddings: The patch embeddings tensor of shape
+                (batch_size, hidden_size, n_patches_height, n_patches_width).
+
+        Returns:
+            The logits as a tensor (batch_size, n_classes, n_patches_height,
+            n_patches_width).
+        """
+        batch_size, hidden_size, height, width = patch_embeddings.shape
+        embeddings_reshaped = patch_embeddings.reshape(batch_size, hidden_size, height * width)
+        logits = self._layers(embeddings_reshaped.permute(0, 2, 1))
+        return logits.permute(0, 2, 1).reshape(batch_size, -1, height, width)
+
+    def _cls_seg(
+        self,
+        logits: torch.Tensor,
+        image_size: Tuple[int, int],
+    ) -> torch.Tensor:
+        """Classify each pixel of the image.
+
+        Args:
+            logits: The decoder outputs of shape (batch_size, n_classes,
+                height, width).
+            image_size: The target image size (height, width).
+
+        Returns:
+            Tensor containing scores for all of the classes with shape
+            (batch_size, n_classes, image_height, image_width).
+        """
+        return functional.interpolate(logits, image_size, mode="bilinear")
+
+    def forward(
+        self,
+        features: List[torch.Tensor],
+        image_size: Tuple[int, int],
+    ) -> torch.Tensor:
+        """Maps the patch embeddings to a segmentation mask of the image size.
+
+        Args:
+            features: List of multi-level image features of shape (batch_size,
+                hidden_size, n_patches_height, n_patches_width).
+            image_size: The target image size (height, width).
+
+        Returns:
+            Tensor containing scores for all of the classes with shape
+            (batch_size, n_classes, image_height, image_width).
+        """
+        patch_embeddings = self._forward_features(features)
+        logits = self._forward_head(patch_embeddings)
+        return self._cls_seg(logits, image_size)

eva/vision/models/wrappers/__init__.py

@@ -0,0 +1,6 @@
+"""Vision Model Wrappers API."""
+
+from eva.vision.models.wrappers.from_registry import ModelFromRegistry
+from eva.vision.models.wrappers.from_timm import TimmModel
+
+__all__ = ["TimmModel", "ModelFromRegistry"]

eva/vision/models/wrappers/from_registry.py

@@ -0,0 +1,48 @@
+"""Vision backbone helper class."""
+
+from typing import Any, Callable, Dict
+
+from typing_extensions import override
+
+from eva.core.models import wrappers
+from eva.vision.models.networks.backbones import BackboneModelRegistry
+
+
+class ModelFromRegistry(wrappers.BaseModel):
+    """Wrapper class for vision backbone models.
+
+    This class can be used by load backbones available in eva's
+    model registry by name. New backbones can be registered by using
+    the `@register_model(model_name)` decorator.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        model_kwargs: Dict[str, Any] | None = None,
+        model_extra_kwargs: Dict[str, Any] | None = None,
+        tensor_transforms: Callable | None = None,
+    ) -> None:
+        """Initializes the model.
+
+        Args:
+            model_name: The name of the model to load.
+            model_kwargs: The arguments used for instantiating the model.
+            model_extra_kwargs: Extra arguments used for instantiating the model.
+            tensor_transforms: The transforms to apply to the output tensor
+                produced by the model.
+        """
+        super().__init__(tensor_transforms=tensor_transforms)
+
+        self._model_name = model_name
+        self._model_kwargs = model_kwargs or {}
+        self._model_extra_kwargs = model_extra_kwargs or {}
+
+        self.load_model()
+
+    @override
+    def load_model(self) -> None:
+        self._model = BackboneModelRegistry.load_model(
+            self._model_name, self._model_kwargs | self._model_extra_kwargs
+        )
+        ModelFromRegistry.__name__ = self._model_name