File size: 1,914 Bytes
0c301a6 b5aeba7 0c301a6 03c4550 0c301a6 |
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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
from PIL import Image
import torch
import timm
import torchvision
import torchvision.transforms as T
from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
import gradio as gr
torch.set_grad_enabled(False);
with open("imagenet_classes.txt", "r") as f:
imagenet_categories = [s.strip() for s in f.readlines()]
transform = T.Compose([
T.Resize(256, interpolation=3),
T.CenterCrop(224),
T.ToTensor(),
T.Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD),
])
model = torch.hub.load('facebookresearch/deit:main', 'deit_base_patch16_224', pretrained=True)
def detr(im):
img = transform(im).unsqueeze(0)
# compute the predictions
out = model(img)
# and convert them into probabilities
scores = torch.nn.functional.softmax(out, dim=-1)[0]
# finally get the index of the prediction with highest score
topk_scores, topk_label = torch.topk(scores, k=5, dim=-1)
d = {}
for i in range(5):
pred_name = imagenet_categories[topk_label[i]]
pred_name = f"{pred_name:<25}"
score = topk_scores[i]
score = f"{score:.3f}"
d[pred_name] = score
return d
inputs = gr.inputs.Image(type='pil', label="Original Image")
outputs = gr.outputs.Label(type="confidences",num_top_classes=5)
title = "Deit"
description = "Gradio demo for Facebook DeiT: Data-efficient Image Transformers. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2012.12877' target='_blank'>Training data-efficient image transformers & distillation through attention</a> | <a href='https://github.com/facebookresearch/deit' target='_blank'>Github Repo</a></p>"
examples = [
['deer.jpeg'],
['cat.jpeg']
]
gr.Interface(detr, inputs, outputs, title=title, description=description, article=article, examples=examples).launch() |