Update handler.py

handler.py

@@ -2,7 +2,8 @@ from typing import Dict, List, Any
 from transformers import pipeline
 import transformers
 from PIL import Image
-import requests
+import base64
+from io import BytesIO
 
 print('TRANSFORMERS VERSION')
 print(transformers.__version__)
@@ -14,12 +15,15 @@ class EndpointHandler():
 
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        # get inputs
-        # image = data.pop("image",data)
+        base64_image = data.get("image")
+        if base64_image is None:
+            raise ValueError("No image provided")
 
-        # load image
-        url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
-        image = Image.open(requests.get(url, stream=True).raw)
+        if base64_image.startswith('data:image/jpeg;base64,'):
+            base64_image = base64_image.replace('data:image/jpeg;base64,', '')
+
+        image_bytes = base64.b64decode(base64_image)
+        image = Image.open(BytesIO(image_bytes))
 
         depth = self.pipe(image)["depth"]