Spaces:
Running
Running
File size: 2,210 Bytes
ec06154 | 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 | import torch
from transformers import DetrForObjectDetection, DetrImageProcessor
from PIL import Image
import gradio as gr
import matplotlib.pyplot as plt
import matplotlib.patches as patches
# Load pre-trained model and processor
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
def detect_car(image: Image.Image) -> Image.Image:
# Preprocess the input image
inputs = processor(images=image, return_tensors="pt")
# Run the model to get predictions
outputs = model(**inputs)
# Postprocess the outputs to get bounding boxes and labels
target_sizes = torch.tensor([image.size[::-1]]) # (height, width)
results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)[0]
# Plotting the image with bounding boxes for objects
fig, ax = plt.subplots(1, figsize=(12, 8))
ax.imshow(image)
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
if score > 0.7: # Confidence threshold for detecting cars
xmin, ymin, xmax, ymax = box.detach().numpy()
width, height = xmax - xmin, ymax - ymin
rect = patches.Rectangle((xmin, ymin), width, height, linewidth=2, edgecolor='red', facecolor='none')
ax.add_patch(rect)
ax.text(xmin, ymin, f"{model.config.id2label[label.item()]}: {score:.2f}",
color='white', fontsize=12, bbox=dict(facecolor='red', alpha=0.5))
# Convert the plot to an image
plt.axis('off')
plt.tight_layout()
# Save the figure to a canvas and convert to image
fig.canvas.draw()
result_img = Image.frombytes('RGB', fig.canvas.get_width_height(), fig.canvas.tostring_rgb())
plt.close(fig)
return result_img
# Gradio interface to upload images and get object detection results
iface = gr.Interface(
fn=detect_car,
inputs=gr.Image(type="pil"),
outputs=gr.Image(type="pil"),
title="Car Detection with DETR",
description="Upload an image and the model will detect cars with bounding boxes. Only cars will be displayed."
)
if __name__ == "__main__":
iface.launch()
|