|
from PIL import Image |
|
import requests |
|
import matplotlib.pyplot as plt |
|
import torch |
|
from torch import nn |
|
from torchvision.models import resnet50 |
|
import torchvision.transforms as T |
|
torch.set_grad_enabled(False); |
|
import gradio as gr |
|
import io |
|
|
|
model = torch.hub.load('facebookresearch/detr', 'detr_resnet50', pretrained=True) |
|
|
|
|
|
torch.hub.download_url_to_file('https://images.pexels.com/photos/461717/pexels-photo-461717.jpeg', 'horse.jpeg') |
|
torch.hub.download_url_to_file('https://images.pexels.com/photos/5967799/pexels-photo-5967799.jpeg', 'turtle.jpeg') |
|
|
|
|
|
|
|
CLASSES = [ |
|
'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', |
|
'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', |
|
'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', |
|
'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', |
|
'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', |
|
'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', |
|
'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', |
|
'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', |
|
'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', |
|
'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A', |
|
'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', |
|
'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', |
|
'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', |
|
'toothbrush' |
|
] |
|
|
|
|
|
COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125], |
|
[0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]] |
|
|
|
|
|
transform = T.Compose([ |
|
T.Resize(800), |
|
T.ToTensor(), |
|
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) |
|
]) |
|
|
|
|
|
def box_cxcywh_to_xyxy(x): |
|
x_c, y_c, w, h = x.unbind(1) |
|
b = [(x_c - 0.5 * w), (y_c - 0.5 * h), |
|
(x_c + 0.5 * w), (y_c + 0.5 * h)] |
|
return torch.stack(b, dim=1) |
|
|
|
def rescale_bboxes(out_bbox, size): |
|
img_w, img_h = size |
|
b = box_cxcywh_to_xyxy(out_bbox) |
|
b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32) |
|
return b |
|
|
|
def fig2img(fig): |
|
"""Convert a Matplotlib figure to a PIL Image and return it""" |
|
buf = io.BytesIO() |
|
fig.savefig(buf) |
|
buf.seek(0) |
|
return Image.open(buf) |
|
|
|
|
|
def plot_results(pil_img, prob, boxes): |
|
plt.figure(figsize=(16,10)) |
|
plt.imshow(pil_img) |
|
ax = plt.gca() |
|
colors = COLORS * 100 |
|
for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors): |
|
ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, |
|
fill=False, color=c, linewidth=3)) |
|
cl = p.argmax() |
|
text = f'{CLASSES[cl]}: {p[cl]:0.2f}' |
|
ax.text(xmin, ymin, text, fontsize=15, |
|
bbox=dict(facecolor='yellow', alpha=0.5)) |
|
plt.axis('off') |
|
return fig2img(plt) |
|
|
|
|
|
|
|
def detr(im): |
|
|
|
img = transform(im).unsqueeze(0) |
|
|
|
|
|
outputs = model(img) |
|
|
|
|
|
probas = outputs['pred_logits'].softmax(-1)[0, :, :-1] |
|
keep = probas.max(-1).values > 0.9 |
|
|
|
|
|
bboxes_scaled = rescale_bboxes(outputs['pred_boxes'][0, keep], im.size) |
|
return plot_results(im, probas[keep], bboxes_scaled) |
|
|
|
|
|
|
|
inputs = gr.inputs.Image(type='pil', label="Original Image", shape=(600,600)) |
|
outputs = gr.outputs.Image(type="pil",label="Output Image") |
|
|
|
examples = [ |
|
['horse.jpeg'], |
|
['turtle.jpeg'] |
|
] |
|
|
|
title = "DETR" |
|
description = "Gradio demo for Facebook DETR. 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/2005.12872'>End-to-End Object Detection with Transformers</a> | <a href='https://github.com/facebookresearch/detr'>Github Repo</a></p>" |
|
|
|
gr.Interface(detr, inputs, outputs, title=title, description=description, article=article, examples=examples).launch() |
|
|
