generated from bananaml/serverless-template
-
Notifications
You must be signed in to change notification settings - Fork 0
/
app.py
38 lines (30 loc) · 1.02 KB
/
app.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import torch
from model.model import OnnxModel
from PIL import Image
import time
import requests
from io import BytesIO
# Init is ran on server startup
# Load your model to GPU as a global variable here using the variable name "model"
def init():
global model
device = 0 if torch.cuda.is_available() else -1
global model
path = "model/classifier.onnx"
model = OnnxModel(path)
# Inference is ran for every server call
# Reference your preloaded global model variable here.
def inference(model_inputs:dict) -> dict:
start_time = time.time()
global model
# Parse out your arguments
img_path = model_inputs.get('input', None)
response = requests.get(img_path)
img = Image.open(BytesIO(response.content))
if img == None:
return {'message': "No image provided"}
# Run the model
result = model.predict(img)
end_time = time.time() - start_time
# Return the results as a dictionary
return {'output' : torch.argmax(torch.Tensor(result)).tolist(), 'time' : end_time}