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

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

eva/vision/models/wrappers/from_timm.py

@@ -0,0 +1,68 @@
+"""Model wrapper for timm models."""
+
+from typing import Any, Callable, Dict, Tuple
+from urllib import parse
+
+import timm
+from typing_extensions import override
+
+from eva.core.models import wrappers
+
+
+class TimmModel(wrappers.BaseModel):
+    """Model wrapper for `timm` models.
+
+    Note that only models with `forward_intermediates`
+    method are currently supported.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        pretrained: bool = True,
+        checkpoint_path: str = "",
+        out_indices: int | Tuple[int, ...] | None = None,
+        model_kwargs: Dict[str, Any] | None = None,
+        tensor_transforms: Callable | None = None,
+    ) -> None:
+        """Initializes the encoder.
+
+        Args:
+            model_name: Name of model to instantiate.
+            pretrained: If set to `True`, load pretrained ImageNet-1k weights.
+            checkpoint_path: Path of checkpoint to load.
+            out_indices: Returns last n blocks if `int`, all if `None`, select
+                matching indices if sequence.
+            model_kwargs: Extra model arguments.
+            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._pretrained = pretrained
+        self._checkpoint_path = checkpoint_path
+        self._out_indices = out_indices
+        self._model_kwargs = model_kwargs or {}
+
+        self.load_model()
+
+    @override
+    def load_model(self) -> None:
+        """Builds and loads the timm model as feature extractor."""
+        self._model = timm.create_model(
+            model_name=self._model_name,
+            pretrained=True if self._checkpoint_path else self._pretrained,
+            pretrained_cfg=self._pretrained_cfg,
+            out_indices=self._out_indices,
+            features_only=self._out_indices is not None,
+            **self._model_kwargs,
+        )
+        TimmModel.__name__ = self._model_name
+
+    @property
+    def _pretrained_cfg(self) -> Dict[str, Any]:
+        if not self._checkpoint_path:
+            return {}
+        key = "file" if parse.urlparse(self._checkpoint_path).scheme in ("file", "") else "url"
+        return {key: self._checkpoint_path, "num_classes": 0}

eva/vision/utils/colormap.py

@@ -0,0 +1,77 @@
+"""Color mapping constants."""
+
+COLORS = [
+    (0, 0, 0),
+    (255, 0, 0),  # Red
+    (0, 255, 0),  # Green
+    (0, 0, 255),  # Blue
+    (255, 255, 0),  # Yellow
+    (255, 0, 255),  # Magenta
+    (0, 255, 255),  # Cyan
+    (128, 128, 0),  # Olive
+    (128, 0, 128),  # Purple
+    (0, 128, 128),  # Teal
+    (192, 192, 192),  # Silver
+    (128, 128, 128),  # Gray
+    (255, 165, 0),  # Orange
+    (210, 105, 30),  # Chocolate
+    (0, 128, 0),  # Lime
+    (255, 192, 203),  # Pink
+    (255, 69, 0),  # Red-Orange
+    (255, 140, 0),  # Dark Orange
+    (0, 255, 255),  # Sky Blue
+    (0, 255, 127),  # Spring Green
+    (0, 0, 139),  # Dark Blue
+    (255, 20, 147),  # Deep Pink
+    (139, 69, 19),  # Saddle Brown
+    (0, 100, 0),  # Dark Green
+    (106, 90, 205),  # Slate Blue
+    (138, 43, 226),  # Blue-Violet
+    (218, 165, 32),  # Goldenrod
+    (199, 21, 133),  # Medium Violet Red
+    (70, 130, 180),  # Steel Blue
+    (165, 42, 42),  # Brown
+    (128, 0, 0),  # Maroon
+    (255, 0, 255),  # Fuchsia
+    (210, 180, 140),  # Tan
+    (0, 0, 128),  # Navy
+    (139, 0, 139),  # Dark Magenta
+    (144, 238, 144),  # Light Green
+    (46, 139, 87),  # Sea Green
+    (255, 255, 0),  # Gold
+    (154, 205, 50),  # Yellow Green
+    (0, 191, 255),  # Deep Sky Blue
+    (0, 250, 154),  # Medium Spring Green
+    (250, 128, 114),  # Salmon
+    (255, 105, 180),  # Hot Pink
+    (204, 255, 204),  # Pastel Light Green
+    (51, 0, 51),  # Very Dark Magenta
+    (255, 102, 0),  # Dark Orange
+    (0, 255, 0),  # Bright Green
+    (51, 153, 255),  # Blue-Purple
+    (51, 51, 255),  # Bright Blue
+    (204, 0, 0),  # Dark Red
+    (90, 90, 90),  # Very Dark Gray
+    (255, 255, 51),  # Pastel Yellow
+    (255, 153, 255),  # Pink-Magenta
+    (153, 0, 76),  # Dark Pink
+    (51, 25, 0),  # Very Dark Brown
+    (102, 51, 0),  # Dark Brown
+    (0, 0, 51),  # Very Dark Blue
+    (180, 180, 180),  # Dark Gray
+    (102, 255, 204),  # Pastel Green
+    (0, 102, 0),  # Dark Green
+    (220, 245, 20),  # Lime Yellow
+    (255, 204, 204),  # Pastel Pink
+    (0, 204, 255),  # Pastel Blue
+    (240, 240, 240),  # Light Gray
+    (153, 153, 0),  # Dark Yellow
+    (102, 0, 51),  # Dark Red-Pink
+    (0, 51, 0),  # Very Dark Green
+    (255, 102, 204),  # Magenta Pink
+    (204, 0, 102),  # Red-Pink
+]
+"""RGB colors."""
+
+COLORMAP = dict(enumerate(COLORS)) | {255: (255, 255, 255)}
+"""Class id to RGB color mapping."""

eva/vision/utils/convert.py

@@ -1,24 +1,67 @@
 """Image conversion related functionalities."""
 
-from typing import Any
+from typing import Iterable
 
-import numpy as np
-import numpy.typing as npt
+import torch
+from torchvision.transforms.v2 import functional
 
 
-def to_8bit(image_array: npt.NDArray[Any]) -> npt.NDArray[np.uint8]:
-    """Casts an image of higher bit image (i.e. 16bit) to 8bit.
+def descale_and_denorm_image(
+    image: torch.Tensor,
+    mean: Iterable[float] = (0.0, 0.0, 0.0),
+    std: Iterable[float] = (1.0, 1.0, 1.0),
+    inplace: bool = True,
+) -> torch.Tensor:
+    """De-scales and de-norms an image tensor to (0, 255) range.
 
     Args:
-        image_array: The image array to convert.
+        image: An image float tensor.
+        mean: The mean that the image channels are normalized with.
+        std: The std that the image channels are normalized with.
+        inplace: Whether to perform the operation in-place.
 
     Returns:
-        The image as normalized as a 8-bit format.
+        The image tensor of range (0, 255) range as uint8.
     """
-    if np.issubdtype(image_array.dtype, np.integer):
-        image_array = image_array.astype(np.float64)
+    if not inplace:
+        image = image.clone()
 
-    image_scaled_array = image_array - image_array.min()
-    image_scaled_array /= image_scaled_array.max()
-    image_scaled_array *= 255
-    return image_scaled_array.astype(np.uint8)
+    norm_image = _descale_image(image, mean=mean, std=std)
+    return _denorm_image(norm_image)
+
+
+def _descale_image(
+    image: torch.Tensor,
+    mean: Iterable[float] = (0.0, 0.0, 0.0),
+    std: Iterable[float] = (1.0, 1.0, 1.0),
+) -> torch.Tensor:
+    """De-scales an image tensor to (0., 1.) range.
+
+    Args:
+        image: An image float tensor.
+        mean: The normalized channels mean values.
+        std: The normalized channels std values.
+
+    Returns:
+        The de-normalized image tensor of range (0., 1.).
+    """
+    return functional.normalize(
+        image,
+        mean=[-cmean / cstd for cmean, cstd in zip(mean, std, strict=False)],
+        std=[1 / cstd for cstd in std],
+    )
+
+
+def _denorm_image(image: torch.Tensor) -> torch.Tensor:
+    """De-normalizes an image tensor from (0., 1.) to (0, 255) range.
+
+    Args:
+        image: An image float tensor.
+
+    Returns:
+        The image tensor of range (0, 255) range as uint8.
+    """
+    image_scaled = image - image.min()
+    image_scaled /= image_scaled.max()
+    image_scaled *= 255
+    return image_scaled.to(dtype=torch.uint8)

eva/vision/utils/io/__init__.py

@@ -1,12 +1,18 @@
 """Vision I/O utilities."""
 
-from eva.vision.utils.io.image import read_image
-from eva.vision.utils.io.nifti import fetch_total_nifti_slices, read_nifti_slice
+from eva.vision.utils.io.image import read_image, read_image_as_array, read_image_as_tensor
+from eva.vision.utils.io.mat import read_mat, save_mat
+from eva.vision.utils.io.nifti import fetch_nifti_shape, read_nifti, save_array_as_nifti
 from eva.vision.utils.io.text import read_csv
 
 __all__ = [
     "read_image",
-    "fetch_total_nifti_slices",
-    "read_nifti_slice",
+    "read_image_as_array",
+    "read_image_as_tensor",
+    "fetch_nifti_shape",
+    "read_nifti",
+    "save_array_as_nifti",
     "read_csv",
+    "read_mat",
+    "save_mat",
 ]

eva/vision/utils/io/image.py

@@ -3,6 +3,8 @@
 import cv2
 import numpy as np
 import numpy.typing as npt
+from torchvision import tv_tensors
+from torchvision.transforms.v2 import functional
 
 from eva.vision.utils.io import _utils
 
@@ -14,7 +16,7 @@ def read_image(path: str) -> npt.NDArray[np.uint8]:
         path: The path of the image file.
 
     Returns:
-        The RGB image as a numpy array.
+        The RGB image as a numpy array (HxWxC).
 
     Raises:
         FileExistsError: If the path does not exist or it is unreachable.
@@ -23,6 +25,23 @@ def read_image(path: str) -> npt.NDArray[np.uint8]:
     return read_image_as_array(path, cv2.IMREAD_COLOR)
 
 
+def read_image_as_tensor(path: str) -> tv_tensors.Image:
+    """Reads and loads the image from a file path as a RGB torch tensor.
+
+    Args:
+        path: The path of the image file.
+
+    Returns:
+        The RGB image as a torch tensor (CxHxW).
+
+    Raises:
+        FileExistsError: If the path does not exist or it is unreachable.
+        IOError: If the image could not be loaded.
+    """
+    image_array = read_image(path)
+    return functional.to_image(image_array)
+
+
 def read_image_as_array(path: str, flags: int = cv2.IMREAD_UNCHANGED) -> npt.NDArray[np.uint8]:
     """Reads and loads an image file as a numpy array.
 
@@ -51,4 +70,4 @@ def read_image_as_array(path: str, flags: int = cv2.IMREAD_UNCHANGED) -> npt.NDA
     if image.ndim == 2 and flags == cv2.IMREAD_COLOR:
         image = image[:, :, np.newaxis]
 
-    return np.asarray(image).astype(np.uint8)
+    return np.asarray(image, dtype=np.uint8)

eva/vision/utils/io/mat.py

@@ -0,0 +1,36 @@
+"""mat I/O related functions."""
+
+import os
+from typing import Any, Dict
+
+import numpy.typing as npt
+import scipy.io
+
+from eva.vision.utils.io import _utils
+
+
+def read_mat(path: str) -> Dict[str, npt.NDArray[Any]]:
+    """Reads and loads a mat file.
+
+    Args:
+        path: The path to the mat file.
+
+    Returns:
+        mat file as dictionary.
+
+    Raises:
+        FileExistsError: If the path does not exist or it is unreachable.
+    """
+    _utils.check_file(path)
+    return scipy.io.loadmat(path)
+
+
+def save_mat(path: str, data: Dict[str, npt.NDArray[Any]]) -> None:
+    """Saves a mat file.
+
+    Args:
+        path: The path to save the mat file.
+        data: The dictionary containing the data to save.
+    """
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    scipy.io.savemat(path, data)