Spaces:
Runtime error
Runtime error
File size: 1,833 Bytes
ac1c6ae cfe24db ac1c6ae 6d93dbe ac1c6ae 564548c ac1c6ae 6d93dbe ac1c6ae 6d93dbe cfe24db fbfbcf2 6d93dbe fbfbcf2 7db538c b56a778 |
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 |
from transformers import AutoProcessor, CLIPModel
import torch
class CLIPImageEncoder:
"""
A class for encoding images using the CLIP model.
Args:
device (str): The device to run the model on (default: "cpu").
Attributes:
device (str): The device to run the model on.
model (CLIPModel): The CLIP model used for image encoding.
processor (AutoProcessor): The tokenizer and input processor for the CLIP model.
"""
def __init__(self, device="cpu"):
self.device = device
self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
self.processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32")
def encode_image(self, image_pil):
"""
Encodes a single image using the CLIP model.
Args:
image_pil: A PIL Image object representing the image to encode.
Returns:
numpy.ndarray: The CLIP embedding for the image.
"""
with torch.no_grad():
input = self.processor(images=image_pil, return_tensors="pt")
image_features = self.model.get_image_features(**input)
return image_features.cpu().detach().numpy()[0]
def encode_images(self, batch):
"""
Encodes a batch of images using the CLIP model.
Args:
batch (Dict[str, Any]): A dictionary containing the batch of images to encode.
Returns:
Dict[str, Any]: A dictionary containing the CLIP embeddings for the batch of images.
"""
images = batch["image"]
input = self.processor(images=images, return_tensors="pt")
with torch.no_grad():
image_features = self.model.get_image_features(**input)
return {"clip_embeddings": image_features.cpu().detach().numpy()}
|