initial commit

app.py

@@ -0,0 +1,60 @@
+import gradio as gr
+import matplotlib.pyplot as plt
+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. 
+"""
+
+image_processor = AutoImageProcessor.from_pretrained(
+    "devonho/detr-resnet-50_finetuned_cppe5"
+)
+model = AutoModelForObjectDetection.from_pretrained(
+    "devonho/detr-resnet-50_finetuned_cppe5"
+)
+
+# Gradio Components
+
+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()