File size: 1,238 Bytes
8ef501c | 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 | import gradio as gr
from PIL import ImageDraw
from transformers import pipeline
object_detector = None
def get_detector():
global object_detector
if object_detector is None:
object_detector = pipeline("object-detection", model="facebook/detr-resnet-50")
return object_detector
def draw_bounding_boxes(image, object_detections):
draw = ImageDraw.Draw(image)
for detection in object_detections:
box = detection["box"]
label = detection["label"]
score = detection["score"]
draw.rectangle(
(box["xmin"], box["ymin"], box["xmax"], box["ymax"]),
outline=(255, 0, 0),
width=2,
)
draw.text((box["xmin"], box["ymin"] - 20), f"{label} ({score:.2f})", fill=(255, 0, 0))
return image
def detect_object(image):
output = get_detector()(image)
return draw_bounding_boxes(image, output)
demo = gr.Interface(
fn=detect_object,
inputs=[gr.Image(label="Select Image", type="pil")],
outputs=[gr.Image(label="Processed Image", type="pil")],
title="@IT AI Enthusiast - Project 6: Object Detector",
description="Detect objects in the provided input image.",
)
demo.launch()
|