Update handler.py

handler.py

@@ -1,37 +1,37 @@
-from typing import Dict, List, Any
+from transformers import CLIPProcessor, CLIPModel, pipeline
 from PIL import Image
+import requests
 from io import BytesIO
 import base64
-import torch
-import open_clip
+from typing import Dict, List, Any
 
 class EndpointHandler():
     def __init__(self, path=""):
-        self.model, self.preprocess, _ = open_clip.create_model_and_transforms('hf-hub:laion/CLIP-ViT-H-14-laion2B-s32B-b79K')
-        self.tokenizer = open_clip.get_tokenizer('hf-hub:laion/CLIP-ViT-H-14-laion2B-s32B-b79K')
+        # Carrega el model i el processor específics de CLIP
+        self.model = CLIPModel.from_pretrained(path)
+        self.processor = CLIPProcessor.from_pretrained(path)
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        image_base64 = data.get("inputs", None)
-        parameters = data.get("parameters", None)
-        if image_base64 is None or parameters is None:
-            raise ValueError("Input data or parameters not provided")
+        # Crea la pipeline de classificació d'imatges zero-shot
+        self.classifier = pipeline("zero-shot-image-classification", model=self.model, tokenizer=self.processor)
 
-        candidate_labels = parameters.get("candidate_labels", None)
-        if candidate_labels is None:
-            raise ValueError("Candidate labels not provided")
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        image_input = data.get("inputs", None)
+        candidate_labels = data.get("candidate_labels", None)
 
-        image = Image.open(BytesIO(base64.b64decode(image_base64)))
-        image = self.preprocess(image).unsqueeze(0)
-        text = self.tokenizer(candidate_labels)
+        if image_input is None or candidate_labels is None:
+            raise ValueError("Image input or candidate labels not provided")
 
-        with torch.no_grad():
-            image_features = self.model.encode_image(image)
-            text_features = self.model.encode_text(text)
-            image_features /= image_features.norm(dim=-1, keepdim=True)
-            text_features /= text_features.norm(dim=-1, keepdim=True)
+        # Determina si l'input és una URL o una cadena base64
+        if image_input.startswith("http"):
+            response = requests.get(image_input)
+            image = Image.open(BytesIO(response.content))
+        else:
+            # Suposa que l'input és base64 i decodifica-la
+            image_data = base64.b64decode(image_input)
+            image = Image.open(BytesIO(image_data))
 
-            text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
+        # Realitza la classificació zero-shot
+        results = self.classifier(images=image, candidate_labels=candidate_labels)
 
-        results = [{"label": label, "score": score.item()} for label, score in zip(candidate_labels, text_probs[0])]
+        # Torna els resultats processats
         return results
-