ckandemir commited on
Commit
796696b
·
1 Parent(s): 57aeb22

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +8 -6
handler.py CHANGED
@@ -1,3 +1,4 @@
 
1
  from typing import Dict, List, Any
2
  from PIL import Image
3
  import torch
@@ -16,16 +17,17 @@ class EndpointHandler():
16
  self.model = self.model.to(device)
17
 
18
  def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
19
- images = data.get("images")
 
20
  # Check if images is None or empty and handle it appropriately
21
- if not images:
22
- return {"captions": []}
23
 
24
  # Default to "a photography of" if texts not provided
25
- texts = data.get("texts", ["a photography of"] * len(images))
26
 
27
  try:
28
- raw_images = [Image.open(BytesIO(_img)).convert("RGB") for _img in images]
29
  processed_inputs = [
30
  self.processor(img, txt, return_tensors="pt") for img, txt in zip(raw_images, texts)
31
  ]
@@ -43,4 +45,4 @@ class EndpointHandler():
43
  except Exception as e:
44
  # Handle or log the exception and optionally return an error message
45
  print(f"Error during processing: {str(e)}")
46
- return {"error": str(e)}
 
1
+ import requests
2
  from typing import Dict, List, Any
3
  from PIL import Image
4
  import torch
 
17
  self.model = self.model.to(device)
18
 
19
  def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
20
+ image_urls = data.get("images")
21
+
22
  # Check if images is None or empty and handle it appropriately
23
+ if not image_urls:
24
+ return {"captions": [], "error": "No images provided"}
25
 
26
  # Default to "a photography of" if texts not provided
27
+ texts = data.get("texts", ["a photography of"] * len(image_urls))
28
 
29
  try:
30
+ raw_images = [Image.open(requests.get(url, stream=True).raw).convert("RGB") for url in image_urls]
31
  processed_inputs = [
32
  self.processor(img, txt, return_tensors="pt") for img, txt in zip(raw_images, texts)
33
  ]
 
45
  except Exception as e:
46
  # Handle or log the exception and optionally return an error message
47
  print(f"Error during processing: {str(e)}")
48
+ return {"captions": [], "error": str(e)}