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

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

@@ -0,0 +1,67 @@
+"""Image conversion related functionalities."""
+
+from typing import Iterable
+
+import torch
+from torchvision.transforms.v2 import functional
+
+
+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: 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 tensor of range (0, 255) range as uint8.
+    """
+    if not inplace:
+        image = image.clone()
+
+    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)

eva/vision/utils/io/nifti.py

@@ -1,23 +1,27 @@
 """NIfTI I/O related functions."""
 
-from typing import Any
+from typing import Any, Tuple
 
 import nibabel as nib
+import numpy as np
 import numpy.typing as npt
 
 from eva.vision.utils.io import _utils
 
 
-def read_nifti_slice(path: str, slice_index: int) -> npt.NDArray[Any]:
-    """Reads and loads a NIfTI image from a file path as `uint8`.
+def read_nifti(
+    path: str, slice_index: int | None = None, *, use_storage_dtype: bool = True
+) -> npt.NDArray[Any]:
+    """Reads and loads a NIfTI image from a file path.
 
     Args:
         path: The path to the NIfTI file.
-        slice_index: The image slice index to return. If `None`, it will
-            return the full 3D image.
+        slice_index: Whether to read only a slice from the file.
+        use_storage_dtype: Whether to cast the raw image
+            array to the inferred type.
 
     Returns:
-        The image as a numpy array.
+        The image as a numpy array (height, width, channels).
 
     Raises:
         FileExistsError: If the path does not exist or it is unreachable.
@@ -25,20 +29,42 @@ def read_nifti_slice(path: str, slice_index: int) -> npt.NDArray[Any]:
     """
     _utils.check_file(path)
     image_data = nib.load(path)  # type: ignore
-    dtype = image_data.get_data_dtype()  # type: ignore
-    image_slice = image_data.slicer[:, :, slice_index : slice_index + 1]  # type: ignore
-    image_array = image_slice.get_fdata()
-    return image_array.astype(dtype)
+    if slice_index is not None:
+        image_data = image_data.slicer[:, :, slice_index : slice_index + 1]  # type: ignore
 
+    image_array = image_data.get_fdata()  # type: ignore
+    if use_storage_dtype:
+        image_array = image_array.astype(image_data.get_data_dtype())  # type: ignore
 
-def fetch_total_nifti_slices(path: str) -> int:
-    """Fetches the total slides of a NIfTI image file.
+    return image_array
+
+
+def save_array_as_nifti(
+    array: npt.ArrayLike,
+    filename: str,
+    *,
+    dtype: npt.DTypeLike | None = np.int64,
+) -> None:
+    """Saved a numpy array as a NIfTI image file.
+
+    Args:
+        array: The image array to save.
+        filename: The name to save the image like.
+        dtype: The data type to save the image.
+    """
+    nifti_image = nib.Nifti1Image(array, affine=np.eye(4), dtype=dtype)  # type: ignore
+    nifti_image.header.get_xyzt_units()
+    nifti_image.to_filename(filename)
+
+
+def fetch_nifti_shape(path: str) -> Tuple[int]:
+    """Fetches the NIfTI image shape from a file.
 
     Args:
         path: The path to the NIfTI file.
 
     Returns:
-        The number of the total available slides.
+        The image shape.
 
     Raises:
         FileExistsError: If the path does not exist or it is unreachable.
@@ -46,5 +72,4 @@ def fetch_total_nifti_slices(path: str) -> int:
     """
     _utils.check_file(path)
     image = nib.load(path)  # type: ignore
-    image_shape = image.header.get_data_shape()  # type: ignore
-    return image_shape[-1]
+    return image.header.get_data_shape()  # type: ignore

eva/vision/utils/io/text.py

@@ -4,15 +4,22 @@ import csv
 from typing import Dict, List
 
 
-def read_csv(path: str) -> List[Dict[str, str]]:
+def read_csv(
+    path: str,
+    *,
+    delimiter: str = ",",
+    encoding: str = "utf-8",
+) -> List[Dict[str, str]]:
     """Reads a CSV file and returns its contents as a list of dictionaries.
 
     Args:
         path: The path to the CSV file.
+        delimiter: The character that separates fields in the CSV file.
+        encoding: The encoding of the CSV file.
 
     Returns:
         A list of dictionaries representing the data in the CSV file.
     """
-    with open(path, newline="") as file:
-        data = csv.DictReader(file, skipinitialspace=True)
+    with open(path, newline="", encoding=encoding) as file:
+        data = csv.DictReader(file, skipinitialspace=True, delimiter=delimiter)
         return list(data)