import requests
import gradio as gr
import torch

from timm import create_model
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform


IMAGENET_1K_URL = 'https://storage.googleapis.com/bit_models/ilsvrc2012_1k_wordnet_lemmas.txt'
IMAGENET_1K_LABELS = requests.get(IMAGENET_1K_URL).text.strip().split('\n')

model = create_model('resnet50', pretrained=True)

transform = create_transform(
    **resolve_data_config({}, model=model)['test_time_augmentation'][0])


model.eval()


def predict(image):
    img = image.convert('RGB')
    transformed_image = transform(img).unsqueeze(0)
    with torch.no_grad():
        out = model(transformed_image)
    probabilities = torch.nn.functional.softmax(out[0], dim=0)
    values, indices = torch.topk(probabilities, k=5)
    return {IMAGENET_1K_LABELS[i]: v.item() for i, v in zip(indices, values)}


gr.Interface(predict, gr.inputs.Image(type='pil'),
             output='label').launch()