forked from agermanidis/squeezenet
-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathrunway_model.py
37 lines (31 loc) · 918 Bytes
/
runway_model.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
import json
from PIL import Image
from torchvision import models, transforms
from torch.autograd import Variable
import runway
from runway.data_types import image, text
labels = json.load(open('labels.json'))
normalize = transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
preprocess = transforms.Compose([
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize
])
@runway.setup
def setup():
return models.squeezenet1_1(pretrained=True)
@runway.command('classify', inputs={'photo': image() }, outputs={'label': text() })
def classify(model, inputs):
img = inputs['photo']
img_tensor = preprocess(img)
img_tensor.unsqueeze_(0)
img_variable = Variable(img_tensor)
fc_out = model(img_variable)
label = labels[str(fc_out.data.numpy().argmax())]
return {'label': label}
if __name__ == '__main__':
runway.run()