| from typing import Any, Dict, List, AnyStr | |
| import numpy as np | |
| from transformers import CLIPProcessor, CLIPModel | |
| from PIL import Image | |
| from io import BytesIO | |
| import base64 | |
| class EndpointHandler(): | |
| def __init__(self, path="") -> None: | |
| "Preload all the elements we need at inference." | |
| self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") | |
| self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") | |
| self.path = path | |
| def __call__(self, data: Dict[str, AnyStr]) -> List[Dict[str, AnyStr]]: | |
| "Run the inference." | |
| inputs = data.get('inputs') | |
| text = inputs.get('text') | |
| imageData = inputs.get('image') | |
| image = Image.open(BytesIO(base64.b64decode(imageData))) | |
| inputs = self.processor(text=text, images=image, return_tensors="pt", padding=True) | |
| outputs = self.model(**inputs) | |
| image_embeds = outputs.image_embeds.detach().numpy().flatten().tolist() | |
| text_embeds = outputs.text_embeds.detach().numpy().flatten().tolist() | |
| logits_per_image = outputs.logits_per_image.detach().numpy().flatten().tolist() | |
| return {'image_embeddings': image_embeds, 'text_embeddings': text_embeds, 'logits_per_image': logits_per_image} | |