def handler()

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]
    })