Upload image-processing/object-detector.py

image-processing/object-detector.py

@@ -0,0 +1,149 @@
+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()