adds handler.py for custom inference

handler.py

@@ -0,0 +1,94 @@
+import base64
+import io
+import torch
+from typing import Dict, List, Any
+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
+
+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("SuperloreAI/clip-vit-large-patch14")
+        self.processor = CLIPProcessor.from_pretrained("SuperloreAI/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 keys:
+                - "image_list" (List[str]): A list of base64-encoded images.
+                - "text_list" (List[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.
+        """
+        image_list = data.get("image_list", []) # list of b64 images
+        text_list = data.get("text_list", []) # list of texts
+
+        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 & text_features, compute similarity
+        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:
+            # Decode the base64-encoded image and convert it to an RGB image
+            image_data = base64.b64decode(base64_image)
+            image = Image.open(io.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: List[str]) -> torch.Tensor:
+        with torch.no_grad():
+            # Tokenize the input text list
+            input_tokens = self.processor(text_list, return_tensors="pt", padding=True, truncation=True)
+
+            # Generate the embeddings for the text list
+            self.model.eval()
+            text_features = self.model.get_text_features(**input_tokens)
+        return text_features
+        
+