Understanding Real-Time Object Detection with RT-DETR: A Code Deep Dive

Object detection powers everything from smart cameras to autonomous robots. In this post I'll walk through the code behind my RT-DETR Object Detection System, a Python project using the RT-DETR model (PekingU/rtdetr_r50vd) for real-time and batch image processing.

Instead of just running the code, we'll dissect each script to understand how it works, from downloading the model to drawing bounding boxes. Let's explore the magic of computer vision with Python, OpenCV, and Transformers.

What's This Project About?

The RT-DETR Object Detection System uses a transformer-based model to detect objects in webcam feeds or image folders, annotating them with labeled bounding boxes. The project, built with Python, OpenCV, PyTorch, and the Transformers library, consists of three key scripts:

• local_save.py: Downloads and saves the RT-DETR model.
• main.py: Orchestrates detection for webcam or image inputs.
• utils.py: Handles model loading and visualization.

This walkthrough explains what each script does, line by line, to help you understand object detection under the hood.

Full code is on GitHub: DiamondDolby/rtdetr-object-detector

Script 1: local_save.py — Fetching the RT-DETR Model

The local_save.py script prepares the RT-DETR model for use by downloading it from the Hugging Face Transformers library and saving it locally.

from transformers import AutoImageProcessor, AutoModelForObjectDetection

# Pre-Trained RT-DETR model
model_name = "PekingU/rtdetr_r50vd"

# Local Folder
save_path = "src/rtdetr_model"

# Setting up image processor and neural network
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForObjectDetection.from_pretrained(model_name)

# Save to local folder
processor.save_pretrained(save_path)
model.save_pretrained(save_path)

What's happening?

• Imports: We use AutoImageProcessor and AutoModelForObjectDetection from the transformers library to handle image preprocessing and the RT-DETR model.
• Model name: PekingU/rtdetr_r50vd is the pre-trained RT-DETR model hosted on Hugging Face.
• Processor and model: the processor prepares images for the model (resizing, normalizing), while the model is the neural network for detection.
• Saving locally: save_pretrained(save_path) stores the model and processor in src/rtdetr_model for offline use, reducing dependency on internet access.

This script ensures the model is ready before running detection tasks.

Script 2: main.py — The Control Center

The main.py script is the entry point, letting users choose between webcam or image folder modes and coordinating the detection process.

import os
import cv2
import torch
from utils import load_model, draw_boxes

# Load model
processor, model = load_model()
model.eval()

We use cv2 (OpenCV) for image/video handling, torch for model inference, and custom load_model and draw_boxes from utils.py. load_model() retrieves the saved model and processor, and model.eval() sets the model to evaluation mode (disabling training-specific operations).

def detect_from_webcam():
    cap = cv2.VideoCapture(0)  # Change to 1 if using an external webcam
    frame_count = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        frame_count += 1
        if frame_count % 3 != 0:  # change number to skip frames
            continue  # Skip every 3 frames

        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        inputs = processor(images=rgb, return_tensors="pt")

        with torch.no_grad():
            outputs = model(**inputs)

        annotated = draw_boxes(frame, outputs, processor, model)
        cv2.imshow("RT-DETR Live Feed", annotated)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

This is webcam mode (detect_from_webcam). Here's what's happening:

• Opens the default webcam (cv2.VideoCapture(0)).
• Skips every third frame (frame_count % 3) for performance.
• Converts frames to RGB (cv2.cvtColor), since RT-DETR expects RGB input.
• Processes frames with the processor to create model inputs.
• Runs inference (model(**inputs)) without gradient computation for efficiency.
• Calls draw_boxes to annotate frames and displays them (cv2.imshow).
• Exits on pressing q.

def detect_from_images():
    input_folder = "images"
    output_folder = "output"
    os.makedirs(output_folder, exist_ok=True)

    for filename in os.listdir(input_folder):
        if filename.lower().endswith((".jpg", ".png", ".jpeg")):
            path = os.path.join(input_folder, filename)
            image = cv2.imread(path)
            rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            inputs = processor(images=rgb, return_tensors="pt")

            with torch.no_grad():
                outputs = model(**inputs)

            annotated = draw_boxes(image, outputs, processor, model)
            cv2.imwrite(os.path.join(output_folder, filename), annotated)
            print(f"Saved: {filename}")

This is image mode (detect_from_images). It loops through images in the images/ folder, reads and processes each one the same way as webcam mode, then saves annotated images to output/ with cv2.imwrite.

if __name__ == "__main__":
    mode = input("Choose mode (1 = webcam, 2 = image folder): ")
    if mode == "1":
        detect_from_webcam()
    elif mode == "2":
        detect_from_images()
    else:
        print("Invalid mode.")

This is the user interface: the if __name__ == "__main__": block prompts the user to choose a mode via input, making the pipeline user-friendly and versatile.

Script 3: utils.py — The Helper Functions

The utils.py script contains reusable functions for loading the model and visualizing detections.

import cv2
import torch
from transformers import AutoImageProcessor, AutoModelForObjectDetection

def load_model():
    processor = AutoImageProcessor.from_pretrained("src/rtdetr_model")
    model = AutoModelForObjectDetection.from_pretrained("src/rtdetr_model")
    return processor, model

load_model() loads the saved processor and model from src/rtdetr_model, and returns both for use in main.py.

def draw_boxes(image, outputs, processor, model, threshold=0.3):
    target_sizes = [image.shape[:2][::-1]]  # (height, width) -> (width, height)
    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=threshold)[0]
    id2label = model.config.id2label

    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
        box = [int(i) for i in box.tolist()]
        cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
        cv2.putText(image, f"{id2label[label.item()]}: {score:.2f}",
                    (box[0], box[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
    return image

draw_boxes() takes the input image, model outputs, processor, and model. It converts image dimensions to (width, height) for the processor, uses processor.post_process_object_detection to extract detections (boxes, scores, labels) at a confidence threshold of 0.3, maps label IDs to names via model.config.id2label, and draws green rectangles with labels and confidence scores on the image.

Keeping this logic in utils.py separates model loading and visualization from the main execution, so the codebase stays modular.

Why This Matters

The RT-DETR model leverages transformer architecture, using attention mechanisms to focus on relevant image regions, which makes it fast and accurate. Breaking down the code this way shows how it integrates with OpenCV for visualization and PyTorch for inference, a practical introduction to computer vision that applies whether you're processing real-time video or batch images with a state-of-the-art model.

Try It Yourself

The full project is open-source under the MIT License on GitHub: DiamondDolby/rtdetr-object-detector.

To run it: install dependencies (opencv-python, torch, transformers), download the model with local_save.py, and execute main.py. Try experimenting with the confidence threshold, or add new features like video file support.