"""CLIP image embedder built on pyvisim's own CLIP implementation.
The embedder pairs the image towers implemented in :mod:`._model` with
pretrained safetensors weights downloaded from the Hugging Face Hub (see
:mod:`._registry`), so no third-party CLIP library is needed. Every variant
of open_clip's registry whose image tower is a standard CLIP Vision
Transformer or modified ResNet, including all original OpenAI models, is
supported.
"""
from pathlib import Path
from typing import Any, ClassVar, cast
import numpy as np
from PIL import Image
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(package="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.
Matches the transform open_clip builds for the checkpoint: 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 (see :func:`pyvisim.neural_networks.clip.fetch_checkpoint`);
only the image tower is loaded, and it always runs in ``float32``.
:meth:`embed` returns one embedding per image, L2-normalized when
``normalize`` is on, so they can be compared directly with a dot
product or the cosine similarity metric.
Variant names and pretrained tags follow open_clip, e.g.
``ClipEmbedder("ViT-B-32", pretrained="laion2b_s34b_b79k")``. OpenAI-style
variant spellings (``"ViT-B/32"``) are accepted as aliases. See
:func:`pyvisim.neural_networks.clip.available_variants` and
:func:`pyvisim.neural_networks.clip.available_pretrained` for the
supported combinations.
: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, so weights already downloaded via
open_clip's Hub downloads are reused.
: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})"
)