Andrew Luo commited on
Commit
4ac3bf9
1 Parent(s): ec3bbda
Files changed (2) hide show
  1. handler.py +6 -2
  2. requirements.txt +1 -1
handler.py CHANGED
@@ -2,7 +2,7 @@ from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoToken
2
  import torch
3
  from PIL import Image
4
  from typing import Dict, List, Any
5
- import urllib.request
6
 
7
 
8
  class EndpointHandler():
@@ -36,7 +36,11 @@ class EndpointHandler():
36
  image_paths = data.pop("image_paths", data)
37
  images = []
38
  for image_path in image_paths:
39
- urllib.request.urlretrieve(image_path, "temp")
 
 
 
 
40
  i_image = Image.open("temp")
41
  if i_image.mode != "RGB":
42
  i_image = i_image.convert(mode="RGB")
 
2
  import torch
3
  from PIL import Image
4
  from typing import Dict, List, Any
5
+ import requests
6
 
7
 
8
  class EndpointHandler():
 
36
  image_paths = data.pop("image_paths", data)
37
  images = []
38
  for image_path in image_paths:
39
+ response = requests.get(image_path)
40
+ response.raise_for_status() # Raise an exception if the request failed
41
+
42
+ with open("temp", "wb") as f:
43
+ f.write(response.content)
44
  i_image = Image.open("temp")
45
  if i_image.mode != "RGB":
46
  i_image = i_image.convert(mode="RGB")
requirements.txt CHANGED
@@ -1,4 +1,4 @@
1
  torch
2
  transformers
3
  Pillow
4
- urllib3
 
1
  torch
2
  transformers
3
  Pillow
4
+ requests