Spaces:
Running
Running
| 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() | |