"""CLIP image embedder built on pyvisim's own CLIP implementation."""
from pathlib import Path
from typing import Any, ClassVar, cast
import numpy as np
from PIL import Image
from ..._errors import MissingModuleFromExtraMessage
from ...base import SerializableImageEmbedder
from ...lazy_import import OptionalImport
from ...typing import (
Float32NumpyArray,
UInt8NumpyArray,
)
from ._registry import (
CheckpointSpec,
VisionConfig,
fetch_checkpoint,
get_checkpoint_spec,
get_model_config,
)
with OptionalImport(
err_msg=str(MissingModuleFromExtraMessage(module="torch", extra="nn"))
) as _torch_import:
import torch
from torchvision import transforms
from ...utils.torch_utils import (
decode_state_dict,
encode_state_dict,
resolve_device,
)
from ._model import build_vision_model, load_vision_weights
_torch_import.check()
def _build_preprocess(
config: VisionConfig, spec: CheckpointSpec
) -> "transforms.Compose":
"""
Build the CLIP inference preprocessing pipeline of a checkpoint.
Steps include: A bicubic resize (of the shortest side plus a center crop, or of both sides
for ``"squash"`` checkpoints), conversion to a ``[0, 1]`` tensor and
normalization with the checkpoint's channel statistics.
:param config: Architecture description carrying the input size.
:param spec: Checkpoint metadata carrying the normalization statistics
and resize mode.
:return: The composed torchvision transform.
"""
size = config.image_size
if spec.resize_mode == "squash":
resize: list[transforms.Compose | object] = [
transforms.Resize(
(size, size), interpolation=transforms.InterpolationMode.BICUBIC
)
]
else:
resize = [
transforms.Resize(size, interpolation=transforms.InterpolationMode.BICUBIC),
transforms.CenterCrop(size),
]
return transforms.Compose(
[*resize, transforms.ToTensor(), transforms.Normalize(spec.mean, spec.std)]
)
[docs]
class ClipEmbedder(SerializableImageEmbedder):
"""
Embeds images with a pretrained CLIP model.
The safetensors checkpoint of the requested ``variant`` and
``pretrained`` tag is downloaded from the Hugging Face Hub on first use
and cached.
To view all supported combinations of variant names and pretrained tags,
run:
.. code-block:: python
from pyvisim.neural_networks.clip import available_variants, available_pretrained
for variant in available_variants():
for pretrained in available_pretrained(variant):
print(f"{variant} / {pretrained}")
:param variant: CLIP variant name.
:param pretrained: Pretrained tag naming the weights, e.g. ``"openai"``
for the original OpenAI checkpoint of the variant.
:param device: Device to run the model on (``"cpu"`` or ``"cuda"``).
If ``None``, ``"cuda"`` is used when a CUDA device is available,
else ``"cpu"``. The model runs in ``float32`` on either device.
:param normalize: Whether to L2-normalize the returned embeddings.
:param similarity_func: Name of the built-in similarity metric used to
score two embeddings. One of ``"cosine"``, ``"euclidean"``,
``"l1"`` or ``"manhattan"``.
:param cache_dir: Directory of the Hugging Face Hub cache the checkpoint
is stored in. If ``None``, the standard Hub cache
(``~/.cache/huggingface/hub``) is used.
:param batch_size: Maximum number of images processed in a single batch.
Set to ``-1`` to process all images as a single batch.
:raises ValueError: If ``variant``, ``pretrained`` or
``similarity_func`` is not supported, or the checkpoint does not
match the architecture.
:raises ImportError: If the ``nn`` extra is not installed.
:raises huggingface_hub.errors.HfHubHTTPError: If the checkpoint
download fails.
References:
===========
[1] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel
Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin,
Jack Clark, Gretchen Krueger, and Ilya Sutskever, "Learning
Transferable Visual Models From Natural Language Supervision," in
Proc. ICML, PMLR 139, pp. 8748-8763, 2021.
"""
__format_version__: ClassVar[int] = 4
__state_keys__: ClassVar[frozenset[str]] = (
SerializableImageEmbedder.__state_keys__
| {
"config",
"state_dict",
}
)
def __init__(
self,
variant: str = "ViT-B-32",
pretrained: str = "openai",
*,
device: str | None = None,
normalize: bool = True,
similarity_func: str = "cosine",
cache_dir: str | Path | None = None,
batch_size: int = 16,
) -> None:
super().__init__(
similarity_func=similarity_func,
normalize=normalize,
batch_size=batch_size,
)
self._build(variant, pretrained, device=device)
load_vision_weights(
self._model, fetch_checkpoint(variant, pretrained, cache_dir=cache_dir)
)
def _build(
self,
variant: str,
pretrained: str,
*,
device: str | None,
) -> None:
"""
Build the image tower and the preprocessing of a checkpoint.
The architecture, the preprocessing statistics and the activation
function all follow from the ``variant`` / ``pretrained`` pair, so this
sets up everything but the weights, which the caller loads.
:param variant: CLIP variant name.
:param pretrained: Pretrained tag naming the weights.
:param device: Device to run the model on, or ``None`` to auto-select.
:raises ValueError: If ``variant`` or ``pretrained`` is not supported.
"""
self._config = get_model_config(variant)
self._spec = get_checkpoint_spec(variant, pretrained)
self._variant = variant.replace("/", "-")
self._pretrained = pretrained
self._device = resolve_device(device)
model = build_vision_model(self._config, quick_gelu=self._spec.quick_gelu)
self._model = model.eval().to(self._device)
self._transform = _build_preprocess(self._config, self._spec)
def _state(self) -> dict[str, Any]:
return {
"similarity_func": self._similarity_func_name,
"normalize": self.normalize,
"batch_size": self.batch_size,
"config": {
"variant": self._variant,
"pretrained": self._pretrained,
"device": self._device,
},
"state_dict": encode_state_dict(self._model),
}
[docs]
@classmethod
def from_dict(cls, state: dict[str, Any], **kwargs: Any) -> "ClipEmbedder":
cls._reject_unsupported_kwargs(kwargs)
# __init__ downloads the pretrained checkpoint, which the saved
# state_dict overwrites right after, so the instance is built without it.
embedder = cls.__new__(cls)
SerializableImageEmbedder.__init__(
embedder,
similarity_func=state["similarity_func"],
normalize=state["normalize"],
batch_size=state["batch_size"],
)
config = state["config"]
embedder._build(
config["variant"], config["pretrained"], device=config["device"]
)
embedder._model.load_state_dict(decode_state_dict(state["state_dict"]))
return embedder
@property
def variant(self) -> str:
"""The CLIP variant name in open_clip spelling (e.g. ``"ViT-B-32"``)."""
return self._variant
@property
def pretrained(self) -> str:
"""The pretrained tag naming the loaded weights (e.g. ``"openai"``)."""
return self._pretrained
@property
def device(self) -> str:
return self._device
@property
def embedding_dim(self) -> int:
"""Dimensionality of the embeddings produced by :meth:`embed`."""
return self._config.embed_dim
@property
def image_size(self) -> int:
"""Side length in pixels of the square model input."""
return self._config.image_size
def _preprocess(self, image: UInt8NumpyArray) -> "torch.Tensor":
"""
Convert one canonical ``uint8`` image into a model-ready tensor.
The image is routed through PIL and converted to RGB, so grayscale
``(H, W)`` inputs are accepted as well as ``(H, W, C)`` ones.
:param image: A ``uint8`` array of shape ``(H, W)`` or ``(H, W, C)``.
:return: A normalized image tensor of shape
``(3, image_size, image_size)``.
"""
pil_image = Image.fromarray(image).convert("RGB")
return cast(torch.Tensor, self._transform(pil_image))
@torch.no_grad()
def _embed(self, images: list[UInt8NumpyArray]) -> Float32NumpyArray:
tensors = torch.stack([self._preprocess(image) for image in images])
features = self._model(tensors.to(self._device))
return np.asarray(features.float().cpu().numpy(), dtype=np.float32)
def __repr__(self) -> str:
return (
f"{self.__class__.__name__}(variant={self._variant}, "
f"pretrained={self._pretrained}, device={self._device}, "
f"normalize={self.normalize}, "
f"similarity_func={self._similarity_func_name})"
)