Spaces:
Running
Running
File size: 5,711 Bytes
2231656 | 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | import torch
from transformers import AutoImageProcessor, AutoModelForObjectDetection
from PIL import ImageDraw, ImageFont
import gradio as gr
# ---------------------------
# 1. Load model + processor once
# ---------------------------
# Loading outside the function means we download and initialize
# the model only at startup, not on every prediction.
MODEL_NAME = "facebook/detr-resnet-50"
# Processor: resizes/normalizes the image and converts detection outputs
# back into real pixel coordinates for the original image size.
processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
# Model: DETR (DEtection TRansformer) - detects objects and draws boxes
# around them, along with a predicted label and confidence score.
model = AutoModelForObjectDetection.from_pretrained(MODEL_NAME)
model.eval() # inference mode - disables dropout etc.
# Automatically use GPU if available, otherwise fall back to CPU.
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
# A small color palette to cycle through so each box gets a different color.
COLORS = [
"#e6194b",
"#3cb44b",
"#ffe119",
"#4363d8",
"#f58231",
"#911eb4",
"#46f0f0",
"#f032e6",
"#bcf60c",
"#fabebe",
]
# ---------------------------
# 2. Define the detection logic
# ---------------------------
def detect_objects(image, confidence_threshold: float = 0.7):
"""
Run object detection on the uploaded image and draw bounding boxes
around every object detected above the given confidence threshold.
Args:
image: A PIL image uploaded by the user.
confidence_threshold: Minimum confidence score (0-1) required
to keep a detected object.
Returns:
annotated: The image with bounding boxes drawn on it.
summary: A text list of detected objects and their confidence scores.
"""
# Guard against no image being uploaded.
if image is None:
return None, "Please upload an image."
# Convert the PIL image into the tensor format DETR expects,
# and move it to the same device (CPU/GPU) as the model.
inputs = processor(images=image, return_tensors="pt").to(device)
# Run the model without tracking gradients (faster, less memory - we're not training).
with torch.no_grad():
outputs = model(**inputs)
# DETR predicts boxes in a normalized format; this converts them
# back to actual pixel coordinates based on the original image size.
# image.size is (width, height), but target_sizes expects (height, width).
target_sizes = torch.tensor([image.size[::-1]])
# Keep only detections whose confidence is above the threshold.
results = processor.post_process_object_detection(
outputs, target_sizes=target_sizes, threshold=confidence_threshold
)[0]
# Make a copy of the image so we can draw on it without changing the original.
annotated = image.convert("RGB").copy()
draw = ImageDraw.Draw(annotated)
detections = []
# results["scores"], ["labels"], and ["boxes"] are aligned lists -
# each index i refers to the same detected object.
for i, (score, label_id, box) in enumerate(
zip(results["scores"], results["labels"], results["boxes"])
):
# box is [x_min, y_min, x_max, y_max] in pixel coordinates.
box = [round(coord, 1) for coord in box.tolist()]
# Convert the numeric label ID (e.g. 17) into a readable name (e.g. "cat").
label = model.config.id2label[label_id.item()]
confidence = round(score.item(), 3)
# Cycle through the color palette so each box looks distinct.
color = COLORS[i % len(COLORS)]
# Draw the bounding box rectangle and a text label above it.
draw.rectangle(box, outline=color, width=3)
draw.text((box[0], max(box[1] - 12, 0)), f"{label} ({confidence})", fill=color)
detections.append(f"{label}: {confidence}")
# Build a readable summary string, or a fallback message if nothing was found.
summary = (
"\n".join(detections)
if detections
else "No objects detected above the confidence threshold."
)
return annotated, summary
# ---------------------------
# 3. Build the web UI with Gradio
# ---------------------------
with gr.Blocks(title="Object Detector") as demo:
gr.Markdown("## Object Detector")
gr.Markdown("Powered by Hugging Face DETR (facebook/detr-resnet-50)")
with gr.Row():
# Left side: where the user uploads their image.
image_input = gr.Image(type="pil", label="Upload Image")
# Right side: the same image, but with boxes drawn on it.
image_output = gr.Image(type="pil", label="Detected Objects")
# Lets the user control how strict detection should be.
# Lower threshold = more (possibly less accurate) detections.
confidence_slider = gr.Slider(
minimum=0.1, maximum=0.95, value=0.7, step=0.05, label="Confidence Threshold"
)
detect_btn = gr.Button("Detect Objects")
detections_output = gr.Textbox(label="Detections", lines=6)
# Wire the button click to the detect_objects function:
# takes the image + slider value as input, and updates
# both the annotated image and the text summary.
detect_btn.click(
fn=detect_objects,
inputs=[image_input, confidence_slider],
outputs=[image_output, detections_output],
)
# ---------------------------
# 4. Start the web server
# ---------------------------
if __name__ == "__main__":
demo.launch()
|