rroset commited on
Commit
c2d2a0b
1 Parent(s): 50ba760

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +13 -12
handler.py CHANGED
@@ -11,20 +11,21 @@ class EndpointHandler():
11
  self.classifier = pipeline("zero-shot-image-classification", model="rroset/CLIP-ViT-B-32-laion2B-s34B-b79K")
12
 
13
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
14
- image_input = data.get("inputs", None)
15
- candidate_labels = data.get("candidate_labels", None)
 
16
 
17
- if image_input is None or candidate_labels is None:
18
- raise ValueError("Image input or candidate labels not provided")
19
 
20
- # Carregar la imatge, podria ser via URL o base64
21
- if image_input.startswith("http"):
22
- response = requests.get(image_input)
23
- image = Image.open(BytesIO(response.content))
24
- else:
25
- # Suposa que l'input és base64 i decodifica-la
26
- image_data = base64.b64decode(image_input)
27
- image = Image.open(BytesIO(image_data))
28
 
29
  # Realitza la classificació zero-shot
30
  results = self.classifier(images=image, candidate_labels=candidate_labels)
 
11
  self.classifier = pipeline("zero-shot-image-classification", model="rroset/CLIP-ViT-B-32-laion2B-s34B-b79K")
12
 
13
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
14
+ # Obté l'imatge en base64 i els paràmetres de les dades
15
+ image_base64 = data.get("inputs", None)
16
+ parameters = data.get("parameters", None)
17
 
18
+ if image_base64 is None or parameters is None:
19
+ raise ValueError("Input data or parameters not provided")
20
 
21
+ # Obté les etiquetes candidates dels paràmetres
22
+ candidate_labels = parameters.get("candidate_labels", None)
23
+ if candidate_labels is None:
24
+ raise ValueError("Candidate labels not provided")
25
+
26
+ # Decodifica la imatge des de base64
27
+ image_data = base64.b64decode(image_base64)
28
+ image = Image.open(BytesIO(image_data))
29
 
30
  # Realitza la classificació zero-shot
31
  results = self.classifier(images=image, candidate_labels=candidate_labels)