| import base64 |
| import torch |
| from typing import Dict, List, Any |
| from io import BytesIO |
| from transformers import CLIPProcessor, CLIPModel |
| from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize |
| from PIL import Image |
| from torch.nn.functional import cosine_similarity |
| from typing import Union |
|
|
| max_text_list_length = 30 |
| max_image_list_length = 20 |
|
|
| class EndpointHandler(): |
| def __init__(self, path: str="", image_size: int=224) -> None: |
| """ |
| Initialize the EndpointHandler with a given model path and image size. |
| |
| Args: |
| path (str, optional): Path to the pretrained model. Defaults to an empty string. |
| image_size (int, optional): The size of the images to be processed. Defaults to 224. |
| """ |
| self.model = CLIPModel.from_pretrained("Superlore/clip-vit-large-patch14") |
| self.processor = CLIPProcessor.from_pretrained("Superlore/clip-vit-large-patch14") |
| self.image_transform = Compose([ |
| Resize(image_size, interpolation=3), |
| CenterCrop(image_size), |
| ToTensor(), |
| Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), |
| ]) |
|
|
| def __call__(self, data: Dict[str, Any]) -> Dict[str, list]: |
| """ |
| Process input data containing image and text lists, computing image and text embeddings, |
| and, if both image and text lists are provided, calculate similarity scores between them. |
| |
| Args: |
| data (Dict[str, Any]): A dictionary containing the following key: |
| - "inputs" (Dict[str, list]): A dictionary containing the following keys: |
| - "image_list" (List[str]): A list of base64-encoded images. |
| - "text_list" (Union[List[str], str]): A list of text strings. |
| |
| Returns: |
| Dict[str, list]: A dictionary containing the following keys: |
| - "image_features" (List[List[float]]): A list of image embeddings. |
| - "text_features" (List[List[float]]): A list of text embeddings. |
| - "similarity_scores" (List[List[float]]): A list of similarity scores between image and text embeddings. |
| Empty if either "image_list" or "text_list" is empty. |
| """ |
| if not isinstance(data, dict): |
| raise ValueError("Expected input data to be a dict.") |
| |
| inputs = data.get("inputs", {}) |
|
|
| if not isinstance(inputs, dict): |
| raise ValueError("Expected 'inputs' to be a dict.") |
| |
| image_list = inputs.get("image_list", []) |
| text_list = inputs.get("text_list", []) |
|
|
| if not isinstance(image_list, list): |
| raise ValueError("Expected 'image_list' to be a list.") |
| if not isinstance(text_list, list) and not isinstance(text_list, str): |
| raise ValueError("Expected 'text_list' to be a list or string.") |
| if not all(isinstance(image, str) for image in image_list): |
| raise ValueError("Expected 'image_list' to contain only strings.") |
| if isinstance(text_list, list) and not all(isinstance(text, str) for text in text_list): |
| raise ValueError("Expected 'text_list' to contain only strings.") |
| |
| |
| if isinstance(text_list, str): |
| text_list = [text_list] |
|
|
| if len(image_list) > max_image_list_length: |
| raise ValueError(f"Expected 'image_list' to have a maximum length of {max_image_list_length}.") |
| if len(text_list) > max_text_list_length: |
| raise ValueError(f"Expected 'text_list' to have a maximum length of {max_text_list_length}.") |
| if not all(is_valid_base64_image(image) for image in image_list): |
| raise ValueError("Expected 'image_list' to contain only valid base64-encoded images.") |
| |
| image_features = self.get_image_embeddings(image_list) if len(image_list) > 0 else None |
| text_features = self.get_text_embeddings(text_list) if len(text_list) > 0 else None |
|
|
| result = { |
| "image_features": image_features.tolist() if image_features is not None else [], |
| "text_features": text_features.tolist() if text_features is not None else [], |
| "similarity_scores": [] |
| } |
| |
| if image_features is not None and text_features is not None: |
| similarity_scores = [cosine_similarity(img_feat, text_features) for img_feat in image_features] |
| result["similarity_scores"] = [t.tolist() for t in similarity_scores] |
| |
| return result |
|
|
| def preprocess_images(self, base64_images: List[str]) -> torch.Tensor: |
| """Loads a list of images and applies preprocessing steps.""" |
| preprocessed_images = [] |
| for base64_image in base64_images: |
| |
| image_data = base64.b64decode(base64_image) |
| image = Image.open(BytesIO(image_data)).convert("RGB") |
| preprocessed_image = self.image_transform(image).unsqueeze(0) |
| preprocessed_images.append(preprocessed_image) |
|
|
| return torch.cat(preprocessed_images, dim=0) |
|
|
| def get_image_embeddings(self, base64_images: List[str]) -> torch.Tensor: |
| image_tensors = self.preprocess_images(base64_images) |
|
|
| with torch.no_grad(): |
| self.model.eval() |
| image_features = self.model.get_image_features(pixel_values=image_tensors) |
|
|
| return image_features |
|
|
| def get_text_embeddings(self, text_list: Union[List[str], str]) -> torch.Tensor: |
| with torch.no_grad(): |
| |
| input_tokens = self.processor(text_list, return_tensors="pt", padding=True, truncation=True) |
|
|
| |
| self.model.eval() |
| text_features = self.model.get_text_features(**input_tokens) |
| return text_features |
| |
| |
| def is_valid_base64_image(data: str) -> bool: |
| try: |
| |
| img_data = base64.b64decode(data) |
| |
| |
| img = Image.open(BytesIO(img_data)) |
| |
| |
| img.verify() |
| |
| return True |
| except: |
| return False |
