RF-DETR：实时视觉对象检测开源模型

综合介绍

RF-DETR 是 Roboflow 团队开发的一个开源对象检测模型。它基于 Transformer 架构，核心特点是实时高效。模型在微软 COCO 数据集上首次实现超过 60 AP 的实时检测，同时在 RF100-VL 基准测试中表现突出，适应多种实际场景。它有两种版本：RF-DETR-base（2900 万参数）和 RF-DETR-large（1.28 亿参数）。模型体积小，适合边缘设备部署。代码和预训练权重使用 Apache 2.0 许可，免费开放给社区使用。用户可以从 GitHub 获取资源，轻松训练或部署。

功能列表

• 实时对象检测：在图像或视频中快速识别物体，延迟低。
• 自定义数据集训练：支持用自己的数据调整模型。
• 边缘设备运行：模型轻量，适合资源有限的设备。
• 可调分辨率：用户可以平衡检测速度和精度。
• 预训练模型支持：提供基于 COCO 数据集的预训练权重。
• 视频流处理：能实时分析视频并输出结果。
• ONNX 导出：支持转为 ONNX 格式，方便跨平台部署。
• 多 GPU 训练：可以用多张显卡加速训练过程。

使用帮助

RF-DETR 的使用分为安装、推理和训练三部分。以下是详细步骤，帮助你快速上手。

安装流程

1. 环境准备
   需要 Python 3.9 或以上版本，以及 PyTorch 1.13.0 或更高。如果用 GPU，运行 nvidia-smi 检查驱动。
   • 安装 PyTorch：

     pip install torch>=1.13.0 torchvision>=0.14.0
     

   • 下载代码：

     git clone https://github.com/roboflow/rf-detr.git
     cd rf-detr
     

   • 安装依赖：

     pip install rfdetr
     

     这会自动安装 numpy、supervision 等必要库。

2. 验证安装
   运行以下代码：

   from rfdetr import RFDETRBase
   print("安装成功")

如果没有报错，安装就完成了。

推理操作

RF-DETR 自带 COCO 数据集的预训练模型，可以直接检测图像或视频。

1. 图像检测
   • 示例代码：

     import io
     import requests
     from PIL import Image
     from rfdetr import RFDETRBase
     import supervision as sv
     model = RFDETRBase()
     url = "https://media.roboflow.com/notebooks/examples/dog-2.jpeg"
     image = Image.open(io.BytesIO(requests.get(url).content))
     detections = model.predict(image, threshold=0.5)
     labels = [f"{class_id} {confidence:.2f}" for class_id, confidence in zip(detections.class_id, detections.confidence)]
     annotated_image = image.copy()
     annotated_image = sv.BoxAnnotator().annotate(annotated_image, detections)
     annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections, labels)
     sv.plot_image(annotated_image)
     

   • 这段代码会检测图片中的物体，标注边界框和置信度，然后显示结果。
2. 视频检测
   • 先安装 opencv-python：

     pip install opencv-python
     

   • 示例代码：

     import cv2
     from rfdetr import RFDETRBase
     import supervision as sv
     model = RFDETRBase()
     cap = cv2.VideoCapture("video.mp4")  # 替换为你的视频路径
     while cap.isOpened():
     ret, frame = cap.read()
     if not ret:
     break
     image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
     detections = model.predict(image, threshold=0.5)
     annotated_frame = sv.BoxAnnotator().annotate(frame, detections)
     cv2.imshow("RF-DETR Detection", annotated_frame)
     if cv2.waitKey(1) & 0xFF == ord('q'):
     break
     cap.release()
     cv2.destroyAllWindows()
     

   • 这会逐帧检测视频中的物体并实时显示。
3. 调整分辨率
   • 初始化时可设置分辨率（必须是 56 的倍数）：

     model = RFDETRBase(resolution=560)
     

   • 分辨率越高，精度越好，但速度会变慢。

训练自定义模型

RF-DETR 支持用自己的数据集微调，但数据集需是 COCO 格式，包含 train、valid 和 test 三个子目录。

1. 准备数据集
   • 示例目录结构：

     dataset/
     ├── train/
     │   ├── _annotations.coco.json
     │   ├── image1.jpg
     │   └── image2.jpg
     ├── valid/
     │   ├── _annotations.coco.json
     │   ├── image1.jpg
     │   └── image2.jpg
     └── test/
     ├── _annotations.coco.json
     ├── image1.jpg
     └── image2.jpg
     

   • 可以用 Roboflow 平台生成 COCO 格式数据集：

     from roboflow import Roboflow
     rf = Roboflow(api_key="你的API密钥")
     project = rf.workspace("rf-100-vl").project("mahjong-vtacs-mexax-m4vyu-sjtd")
     dataset = project.version(2).download("coco")
     

2. 开始训练
   • 示例代码：

     from rfdetr import RFDETRBase
     model = RFDETRBase()
     model.train(dataset_dir="./mahjong-vtacs-mexax-m4vyu-sjtd-2", epochs=10, batch_size=4, grad_accum_steps=4, lr=1e-4)
     

   • 训练时，推荐总批次大小（batch_size * grad_accum_steps）为 16。例如，A100 GPU 用 batch_size=16, grad_accum_steps=1；T4 GPU 用 batch_size=4, grad_accum_steps=4。
3. 多 GPU 训练
   • 创建 main.py 文件：

     from rfdetr import RFDETRBase
     model = RFDETRBase()
     model.train(dataset_dir="./dataset", epochs=10, batch_size=4, grad_accum_steps=4, lr=1e-4)
     

   • 在终端运行：

     python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py
     

   • 将 8 替换为你使用的 GPU 数量。注意调整 batch_size 以保持总批次大小稳定。
4. 加载训练结果
   • 训练后会生成两个权重文件：常规权重和 EMA 权重（更稳定）。加载方式：

     model = RFDETRBase(pretrain_weights="./output/model_ema.pt")
     detections = model.predict("image.jpg")
     

ONNX 导出

• 导出为 ONNX 格式，方便在其他平台部署：

  from rfdetr import RFDETRBase
  model = RFDETRBase()
  model.export()
  

• 导出的文件会保存在 output 目录，适合边缘设备优化推理。

应用场景

1. 自动驾驶
   RF-DETR 能实时检测道路上的车辆和行人。它的低延迟和高精度适合嵌入式系统使用。
2. 工业质检
   在工厂流水线上，RF-DETR 可以快速识别零件缺陷。模型轻量，能直接在设备上运行。
3. 视频监控
   RF-DETR 处理监控视频时，能实时检测异常物体或行为。它支持视频流，适合全天候安防。

QA

1. 支持哪些数据集格式？
   只支持 COCO 格式。数据集需包含 train、valid 和 test 子目录，每个目录有对应的 _annotations.coco.json 文件。
2. 如何获取 Roboflow API 密钥？
   登录 https://app.roboflow.com，在账户设置中找到 API 密钥，复制后设置到环境变量 ROBOFLOW_API_KEY。
3. 训练需要多长时间？
   取决于硬件和数据集大小。在 T4 GPU 上，10 个 epoch 可能需要几小时。小数据集用 CPU 也能跑，但速度慢。