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)) 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'), outputs='label').launch()