Spaces:
Sleeping
Sleeping
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) |