File size: 2,410 Bytes
b8d1773 |
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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
from transformers import DetrImageProcessor, DetrForObjectDetection
import torch
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import gradio as gr
import io
# Load the processor and model
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
def detect_and_display_image(image):
# Ensure image is in PIL format
if isinstance(image, bytes):
image = Image.open(io.BytesIO(image))
elif isinstance(image, str):
image = Image.open(image)
# Process the image
inputs = processor(images=image, return_tensors="pt")
# Perform object detection
outputs = model(**inputs)
# Convert outputs to COCO API format
target_sizes = torch.tensor([image.size[::-1]])
results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
# Create a figure and axis for visualization
fig, ax = plt.subplots(1, figsize=(12, 9))
ax.imshow(image)
# Add bounding boxes and labels to the image
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
box = [round(i, 2) for i in box.tolist()]
# Create a Rectangle patch
rect = patches.Rectangle(
(box[0], box[1]),
box[2] - box[0],
box[3] - box[1],
linewidth=2,
edgecolor='red',
facecolor='none'
)
# Add the patch to the Axes
ax.add_patch(rect)
# Add label and confidence score
plt.text(
box[0], box[1],
f'{model.config.id2label[label.item()]}: {round(score.item(), 3)}',
color='red',
fontsize=12,
bbox=dict(facecolor='yellow', alpha=0.5)
)
plt.axis('off') # Hide the axes
# Save the figure to a BytesIO object and return it
buf = io.BytesIO()
plt.savefig(buf, format='png')
buf.seek(0)
return Image.open(buf)
# Create a Gradio interface
iface = gr.Interface(
fn=detect_and_display_image,
inputs=gr.Image(type="pil"),
outputs=gr.Image(type="pil"),
title="Object Detection with DETR",
description="Upload an image to detect objects using the DETR model.",
live=True
)
# Launch the Gradio app
iface.launch()
|