halimb's picture
Create handler.py
41f99fa verified
raw
history blame
No virus
697 Bytes
from typing import Dict, List, Any
from transformers import pipeline
import transformers
from PIL import Image
import requests
print('TRANSFORMERS VERSION')
print(transformers.__version__)
class EndpointHandler():
def __init__(self, path=""):
# load pipe
self.pipe = pipeline(task="depth-estimation", model=path)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
# get inputs
# image = data.pop("image",data)
# load image
url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
image = Image.open(requests.get(url, stream=True).raw)
depth = self.pipe(image)["depth"]
return depth