File size: 3,444 Bytes
33d23e1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | # SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
from typing import List, Optional, Union
import open_clip
import torch
import torch.nn as nn
from PIL import Image
class CLIPEncoder(nn.Module):
"""
A class for encoding images and texts using a specified CLIP model and computing the similarity between them.
Attributes:
-----------
clip_version: str
The version of the CLIP model to be used.
pretrained: str
The pre-trained weights to load.
model: nn.Module
The CLIP model.
preprocess: Callable
The preprocessing transform to apply to the input image.
device: str
The device to which the model is moved.
"""
def __init__(
self,
clip_version: str = "ViT-B/32",
pretrained: Optional[str] = "",
cache_dir: Optional[str] = None,
):
"""
Initializes the CLIPEncoder with the specified CLIP model version and pre-trained weights.
Parameters:
-----------
clip_version: str, optional
The version of the CLIP model to be used. Defaults to 'ViT-B/32'.
pretrained: str, optional
The pre-trained weights to load. If not provided, it defaults based on clip_version.
cache_dir: str, optional
The directory to cache the model. Defaults to None.
device: str, optional
The device to which the model is moved. Defaults to 'cuda'.
"""
super().__init__()
self.clip_version = clip_version
self.pretrained = pretrained if pretrained else self._get_default_pretrained()
self.model, _, self.preprocess = open_clip.create_model_and_transforms(
self.clip_version, pretrained=self.pretrained, cache_dir=cache_dir
)
self.model.eval()
def _get_default_pretrained(self) -> str:
"""Returns the default pretrained weights based on the clip_version."""
return "openai"
@torch.no_grad()
def get_clip_score(self, text: Union[str, List[str]], image: Union[Image.Image, torch.Tensor]) -> torch.Tensor:
"""
Computes the similarity score between the given text(s) and image using the CLIP model.
Parameters:
-----------
text: Union[str, List[str]]
The text or list of texts to compare with the image.
image: Image.Image
The input image.
Returns:
--------
torch.Tensor
The similarity score between the text(s) and image.
"""
# Preprocess the image and move it to the specified device
image = self.preprocess(image).unsqueeze(0)
# Normalize the image features
image_features = self.model.encode_image(image).float()
image_features /= image_features.norm(dim=-1, keepdim=True)
# If a single text string is provided, convert it to a list
if not isinstance(text, (list, tuple)):
text = [text]
# Tokenize the text and move it to the specified device
text = open_clip.tokenize(text)
# Normalize the text features
text_features = self.model.encode_text(text).float()
text_features /= text_features.norm(dim=-1, keepdim=True)
# Compute the similarity between the image and text features
similarity = image_features @ text_features.T
return similarity
|