in pytorch-inference-docker-lambda/app/app.py [0:0]
def handler(event, context):
print('Received event: ' + json.dumps(event, indent=2))
url = event['url']
img = Image.open(urllib.request.urlopen(url))
scaled_img = transform_test(img)
torch_image = scaled_img.unsqueeze(0)
model = torch.jit.load('./resnet34.pt')
predicted_class = model(torch_image).argmax().item()
print('predicted_class: {}'.format(predicted_class))
# Read the categories
with open("imagenet_classes.txt", "r") as f:
categories = [s.strip() for s in f.readlines()]
print('Categories count: {}'.format(len(categories)))
return json.dumps({
"class": categories[predicted_class]
})