File size: 1,338 Bytes
50ba760 ae82746 266aad6 ae82746 266aad6 ae82746 4621849 50ba760 ae82746 266aad6 c2d2a0b ae82746 c2d2a0b ae82746 c2d2a0b ae82746 266aad6 ae82746 266aad6 ae82746 |
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 |
from transformers import pipeline
from PIL import Image
import requests
from io import BytesIO
import base64
from typing import Dict, List, Any
class EndpointHandler():
def __init__(self, path=""):
# Crea la pipeline de classificació d'imatges zero-shot amb el model específic
self.classifier = pipeline("zero-shot-image-classification", model="rroset/CLIP-ViT-B-32-laion2B-s34B-b79K")
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
# Obté l'imatge en base64 i els paràmetres de les dades
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")
# Obté les etiquetes candidates dels paràmetres
candidate_labels = parameters.get("candidate_labels", None)
if candidate_labels is None:
raise ValueError("Candidate labels not provided")
# Decodifica la imatge des de base64
image_data = base64.b64decode(image_base64)
image = Image.open(BytesIO(image_data))
# Realitza la classificació zero-shot
results = self.classifier(images=image, candidate_labels=candidate_labels)
# Torna els resultats processats
return results
|