Object_detector / app.py
someshb07's picture
Create app.py
fa1a7d8 verified
import gradio as gr
from PIL import Image
# Use a pipeline as a high-level helper
from transformers import pipeline
pipe = pipeline("object-detection", model="facebook/detr-resnet-101")
from PIL import Image, ImageDraw, ImageFont
def merge_bboxes(bboxes, threshold=50):
"""
Merges bounding boxes that are close to each other into one.
:param bboxes: List of bounding boxes (xmin, ymin, xmax, ymax)
:param threshold: Maximum distance between boxes to be considered as part of the same object
:return: List of merged bounding boxes
"""
merged_bboxes = []
for bbox in bboxes:
added = False
for i, merged_bbox in enumerate(merged_bboxes):
if (abs(bbox[0] - merged_bbox[0]) < threshold and abs(bbox[1] - merged_bbox[1]) < threshold) or \
(abs(bbox[2] - merged_bbox[2]) < threshold and abs(bbox[3] - merged_bbox[3]) < threshold):
# Merge the boxes by extending the boundaries
new_bbox = (
min(bbox[0], merged_bbox[0]),
min(bbox[1], merged_bbox[1]),
max(bbox[2], merged_bbox[2]),
max(bbox[3], merged_bbox[3])
)
merged_bboxes[i] = new_bbox
added = True
break
if not added:
merged_bboxes.append(bbox)
return merged_bboxes
def draw_bounding_boxes(image, detections):
"""
Draws bounding boxes on the given PIL image based on the object detection results.
Merges close bounding boxes to avoid duplicate detections.
:param image: PIL Image object
:param detections: List of dictionaries containing object detection results
:return: PIL Image with bounding boxes
"""
draw = ImageDraw.Draw(image)
try:
font = ImageFont.truetype("arial.ttf", 40) # Load font (adjust size as needed)
except:
font = ImageFont.load_default() # Use default font if Arial is not available
bboxes = []
labels = []
# Collect all bounding boxes and labels
for obj in detections:
box = obj["box"]
label = f"{obj['label']} ({obj['score']:.2f})"
xmin, ymin, xmax, ymax = box["xmin"], box["ymin"], box["xmax"], box["ymax"]
bboxes.append((xmin, ymin, xmax, ymax))
labels.append(label)
# Merge close bounding boxes
merged_bboxes = merge_bboxes(bboxes)
# Draw bounding boxes and labels for merged boxes
for idx, bbox in enumerate(merged_bboxes):
xmin, ymin, xmax, ymax = bbox
label = labels[idx]
# Draw bounding box
draw.rectangle([xmin, ymin, xmax, ymax], outline="red", width=5)
# Adjust label placement to avoid overlap
label_bbox = draw.textbbox((xmin, ymin), label, font=font)
label_width, label_height = label_bbox[2] - label_bbox[0], label_bbox[3] - label_bbox[1]
label_x = xmin
label_y = ymin - label_height if ymin - label_height > 0 else ymin + 5 # Adjust above or below box
# Draw label background to make text more visible
draw.rectangle([label_x, label_y, label_x + label_width, label_y + label_height], fill="red")
draw.text((label_x, label_y), label, fill="white", font=font)
return image
# Example usage:
# image = Image.open("your_image.jpg")
# detections = [{'score': 0.92, 'label': 'dog', 'box': {'xmin': 2929, 'ymin': 1297, 'xmax': 5067, 'ymax': 3693}}, ...]
# image_with_boxes = draw_bounding_boxes(image, detections)
# image_with_boxes.show()
def detect_object(image):
raw_imagee = image
output=pipe(raw_imagee)
processed_image = draw_bounding_boxes(raw_imagee, output)
return processed_image
demo = gr.Interface(fn = detect_object,
inputs=[gr.Image(label='Please upload the file to detect the data', type='pil')],
outputs=[gr.Image(label='Detected object', type='pil')],
title= 'Object Detector',
description = 'This application will be used to detect image'
)
demo.launch(share=True)