|
import gradio as gr |
|
import numpy as np |
|
import torch |
|
from PIL import Image, ImageDraw |
|
from transformers import AutoImageProcessor, AutoModelForObjectDetection |
|
|
|
description = """ |
|
## This interface is made with 🤗 Gradio. |
|
|
|
Simply upload an image of any person wearning/not-wearing helmet. |
|
""" |
|
model_id = "anindya64/detr-resnet-50-finetuned" |
|
|
|
image_processor = AutoImageProcessor.from_pretrained(model_id) |
|
model = AutoModelForObjectDetection.from_pretrained(model_id) |
|
|
|
|
|
|
|
image_in = gr.components.Image() |
|
image_out = gr.components.Image() |
|
|
|
|
|
def model_inference(img): |
|
with torch.no_grad(): |
|
inputs = image_processor(images=img, return_tensors="pt") |
|
outputs = model(**inputs) |
|
target_sizes = torch.tensor([img.size[::-1]]) |
|
results = image_processor.post_process_object_detection( |
|
outputs, threshold=0.5, target_sizes=target_sizes |
|
)[0] |
|
return results |
|
|
|
|
|
def plot_results(image): |
|
image = Image.fromarray(np.uint8(image)) |
|
results = model_inference(img=image) |
|
draw = ImageDraw.Draw(image) |
|
for score, label, box in zip( |
|
results["scores"], results["labels"], results["boxes"] |
|
): |
|
score = score.item() |
|
box = [round(i, 2) for i in box.tolist()] |
|
x, y, x2, y2 = tuple(box) |
|
draw.rectangle((x, y, x2, y2), outline="red", width=1) |
|
draw.text((x, y), model.config.id2label[label.item()], fill="white") |
|
draw.text((x+0.5, y-0.5), text=str(score), fill='green' if score > 0.7 else 'red') |
|
return image |
|
|
|
|
|
Iface = gr.Interface( |
|
fn=plot_results, |
|
inputs=[image_in], |
|
outputs=image_out, |
|
title="Object Detection Using Fine-Tuned Vision Transformers", |
|
description=description, |
|
).launch() |
|
|
