RF-DETR训练自己的数据集
长期以来，YOLO 系列凭借其出色的速度优势，一直是实时目标检测领域的绝对主导者。然而，2026 年，这一格局正在被一个名为 RF-DETR 的新模型彻底改变。由 Roboflow 开发的 RF-DETR，是首个在 COCO 数据集上突破 60 AP 的实时目标检测模型，其 2XL 版本达到了 60.1 AP50:95 的惊人成绩。这不仅标志着 Transformer 架构在实时检测领域的重大突破，也为整个计算机视觉社区带来了新的选择。RF-DETR模型基于 Transformer 架构，融合 DINOv2 视觉骨干网络，支持目标检测、实例分割和关键点检测任务，提供从 Nano 到 2XLarge 共 6 种尺寸规格。
1. DINOv2
与许多依赖 CNN 的检测器不同，RF-DETR 采用了 DINOv2 作为其视觉骨干网络。DINOv2 是一个基于 Transformer 的自监督学习模型，在多达 1.42 亿张图像上进行过预训练。这使得 RF-DETR 拥有了极其强大的视觉先验知识，能够更好地理解和迁移到新的、未见过的领域（即 domain adaptation）。在衡量模型迁移能力的 RF100-VL 基准测试中，RF-DETR 同样取得了领先地位
2. 神经架构搜索（NAS）
传统的模型系列（如 YOLO）通常采用统一缩放的方式，即对一个大模型进行深度和宽度的缩减来生成小模型，这可能会继承大模型的低效设计。RF-DETR 则引入了 神经架构搜索（NAS） 技术 。它通过训练一个“超网络”（supernet），在训练过程中随机采样不同的内部配置（如层深度、通道数），最终可以为特定的硬件速度限制（如要求在 T4 GPU 上快于 2.5ms）搜索出最优的子网络结构，从而在硬件限制下获得最佳精度 。
3. NMS-Free 设计
这是 RF-DETR 最核心的优势之一。传统的 CNN 检测器（如 YOLOv8、YOLOv9）在推理后，需要使用 非极大值抑制（NMS） 来过滤掉重复的预测框。NMS 不仅增加了计算开销，更重要的是，其计算量会随着图像中物体数量的增加而大幅上升，导致延迟不稳定。
RF-DETR 作为一个端到端的 Transformer 模型，其解码器天生就能输出干净、唯一的预测结果，彻底摆脱了对 NMS 的依赖。这意味着，无论场景是空旷还是拥挤，RF-DETR 都能提供稳定且可预测的推理延迟，这对于自动驾驶、工业质检等要求严苛的实时应用至关重要 。
4. 检测与分割的无缝融合
自 v1.3.0 版本起，RF-DETR 通过统一的 API 原生支持目标检测和实例分割两种任务 。这意味着你可以使用几乎相同的代码来训练检测模型和分割模型，极大地降低了学习和使用成本 。
5. 训练代码
安装rfdetr
pip install rfdetr
RF-DETR 的训练接口 model.train() 支持 COCO JSON 和 YOLO 两种主流数据格式，并会自动识别 。如果你的数据集是 YOLO 格式，确保根目录包含一个 data.yaml 文件，结构如下:
如果想使用coco数据集格式，这里也给出转换代码:
import os import json from tqdm import tqdm import shutil from PIL import Image def yolo_to_coco(label_dir, image_dir, output_dir, classes_txt_path=None, include_neg_samples=True): """ 将 YOLO 格式数据集转换为 COCO 格式，支持负样本（无标注图像），用于 RT-DETR 等模型训练。 Args: label_dir (str): YOLO 标签文件 (.txt) 所在目录。 image_dir (str): 对应图像文件 (.jpg/.png) 所在目录。 output_dir (str): 输出目录，包含图像副本和 _annotations.coco.json。 classes_txt_path (str, optional): classes.txt 路径。若未提供，自动尝试在 label_dir/../classes.txt 查找。 include_neg_samples (bool): 是否包含无标注的负样本图像，默认开启。 """ os.makedirs(output_dir, exist_ok=True) # 自动推断 classes.txt 路径 if classes_txt_path is None: default_classes = os.path.join(os.path.dirname(label_dir.rstrip('/')), 'classes.txt') if os.path.exists(default_classes): classes_txt_path = default_classes else: raise FileNotFoundError("未找到 classes.txt，请通过 classes_txt_path 指定路径。") # 读取类别 with open(classes_txt_path, 'r', encoding='utf-8') as f: class_names = [line.strip() for line in f if line.strip()] categories = [{"id": i, "name": name, "supercategory": "object"} for i, name in enumerate(class_names)] # 获取所有标签文件（.txt） label_files = [f for f in os.listdir(label_dir) if f.endswith('.txt')] if not label_files: raise ValueError(f"在 {label_dir} 中未找到任何 .txt 标签文件！") # 提取已标注的图像basename（用于区分负样本） labeled_basenames = {os.path.splitext(f)[0] for f in label_files} coco_data = { "images": [], "annotations": [], "categories": categories, "info": {"description": "Converted from YOLO to COCO format with negative samples support"}, "licenses": [] } annotation_id = 1 image_id = 1 # 支持的图像扩展名 IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.webp'} # -------------------------- # 第一步：处理有标注的正样本 # -------------------------- print("📌 开始处理有标注的正样本...") for label_file in tqdm(label_files, desc="Processing labeled images"): label_path = os.path.join(label_dir, label_file) base_name = os.path.splitext(label_file)[0] # 查找对应图像 img_file = None for ext in IMG_EXTENSIONS: candidate = base_name + ext if os.path.exists(os.path.join(image_dir, candidate)): img_file = candidate break if img_file is None: print(f"⚠️ 警告：找不到与 {label_file} 对应的图像，跳过。") continue src_img_path = os.path.join(image_dir, img_file) dst_img_path = os.path.join(output_dir, img_file) shutil.copy(src_img_path, dst_img_path) # 使用 PIL 读取图像尺寸 try: with Image.open(dst_img_path) as img: width, height = img.size except Exception as e: print(f"❌ 读取图像 {img_file} 失败: {str(e)}，跳过。") os.remove(dst_img_path) # 清理损坏的图像副本 continue # 添加图像信息 coco_data["images"].append({ "id": image_id, "file_name": img_file, "width": width, "height": height }) # 读取标签并转换 with open(label_path, 'r') as lf: for line in lf: parts = line.strip().split() if len(parts) < 5: continue try: cls_id = int(parts[0]) x_center, y_center, w_norm, h_norm = map(float, parts[1:5]) except ValueError: print(f"⚠️ 标签文件 {label_file} 格式错误，跳过该行。") continue # 转换为绝对坐标 (x_min, y_min, width, height) bbox_width = w_norm * width bbox_height = h_norm * height x_min = (x_center - w_norm / 2) * width y_min = (y_center - h_norm / 2) * height # 确保 bbox 不越界（关键：避免训练时报错） x_min = max(0.0, x_min) y_min = max(0.0, y_min) bbox_width = min(bbox_width, width - x_min) bbox_height = min(bbox_height, height - y_min) area = bbox_width * bbox_height coco_data["annotations"].append({ "id": annotation_id, "image_id": image_id, "category_id": cls_id, "bbox": [x_min, y_min, bbox_width, bbox_height], "area": area, "iscrowd": 0, "segmentation": [] }) annotation_id += 1 image_id += 1 # -------------------------- # 第二步：处理无标注的负样本 # -------------------------- if include_neg_samples: print("\n📌 开始处理无标注的负样本...") # 获取所有图像文件 all_image_files = [f for f in os.listdir(image_dir) if os.path.splitext(f)[1].lower() in IMG_EXTENSIONS] # 筛选出无对应标签的负样本 neg_sample_files = [f for f in all_image_files if os.path.splitext(f)[0] not in labeled_basenames] if not neg_sample_files: print("ℹ️ 未找到无标注的负样本图像。") else: for img_file in tqdm(neg_sample_files, desc="Processing negative samples"): src_img_path = os.path.join(image_dir, img_file) dst_img_path = os.path.join(output_dir, img_file) # 避免重复复制（如果正样本中已有同名文件） if os.path.exists(dst_img_path): print(f"⚠️ 负样本 {img_file} 已存在于输出目录，跳过。") continue shutil.copy(src_img_path, dst_img_path) # 读取图像尺寸 try: with Image.open(dst_img_path) as img: width, height = img.size except Exception as e: print(f"❌ 读取负样本 {img_file} 失败: {str(e)}，跳过。") os.remove(dst_img_path) continue # 仅添加图像信息，无标注 coco_data["images"].append({ "id": image_id, "file_name": img_file, "width": width, "height": height }) image_id += 1 # 保存 COCO JSON output_json = os.path.join(output_dir, "_annotations.coco.json") with open(output_json, 'w', encoding='utf-8') as f: json.dump(coco_data, f, indent=2, ensure_ascii=False) # 统计信息 total_images = len(coco_data["images"]) total_annotations = len(coco_data["annotations"]) neg_samples_count = total_images - len(labeled_basenames) print(f"\n✅ 转换完成！") print(f"📊 统计信息:") print(f" - 总图像数: {total_images} (正样本 {len(labeled_basenames)} + 负样本 {neg_samples_count})") print(f" - 总标注数: {total_annotations}") print(f"📁 输出目录: {output_dir}") print(f"📄 标注文件: {output_json}") if __name__ == "__main__": # === 请根据你的实际情况修改以下路径 === LABEL_DIR = "./test_coco/labels" # YOLO 标签目录（.txt 文件） IMAGE_DIR = "./test_coco/images" # 包含正样本和负样本的图像目录 OUTPUT_DIR = "./coco_with_neg_samples" # 输出目录 CLASSES_TXT = "./classes.txt" # 类别文件（可选） yolo_to_coco( label_dir=LABEL_DIR, image_dir=IMAGE_DIR, output_dir=OUTPUT_DIR, classes_txt_path=CLASSES_TXT, include_neg_samples=True # 开启负样本支持 )
划分数据集:
import json import os import random import shutil from tqdm import tqdm def split_coco_dataset(input_json, image_dir, output_dir, split_ratio=0.8): # 1. 创建输出目录结构 train_dir = os.path.join(output_dir, 'train') val_dir = os.path.join(output_dir, 'valid') os.makedirs(train_dir, exist_ok=True) os.makedirs(val_dir, exist_ok=True) # 2. 读取原始 JSON with open(input_json, 'r', encoding='utf-8') as f: data = json.load(f) images = data['images'] annotations = data['annotations'] categories = data['categories'] # 3. 随机打乱图片顺序 random.seed(42) # 固定随机种子，保证结果可复现 random.shuffle(images) # 4. 计算拆分点 split_idx = int(len(images) * split_ratio) train_images = images[:split_idx] val_images = images[split_idx:] def save_split(split_images, split_name, target_dir): # 提取当前拆分对应的图片 ID 集合 image_ids = {img['id'] for img in split_images} # 过滤对应的标注 split_annos = [anno for anno in annotations if anno['image_id'] in image_ids] # 构造新的 JSON 数据 split_data = { "images": split_images, "annotations": split_annos, "categories": categories, "info": data.get("info", {}), "licenses": data.get("licenses", []) } # 保存 JSON 文件 json_path = os.path.join(target_dir, "_annotations.coco.json") with open(json_path, 'w', encoding='utf-8') as f: json.dump(split_data, f, indent=2, ensure_ascii=False) # 物理复制图片到对应目录 print(f"正在处理 {split_name} 图片...") for img in tqdm(split_images): src_path = os.path.join(image_dir, img['file_name']) dst_path = os.path.join(target_dir, img['file_name']) # 如果图片存在则复制 if os.path.exists(src_path): shutil.copy(src_path, dst_path) else: print(f"⚠️ 警告：找不到图片 {src_path}") # 5. 执行保存逻辑 save_split(train_images, "训练集", train_dir) save_split(val_images, "验证集", val_dir) print(f"\n✅ 划分完成！") print(f"📂 训练集: {len(train_images)} 张图，存放在 {train_dir}") print(f"📂 验证集: {len(val_images)} 张图，存放在 {val_dir}") if __name__ == "__main__": # === 请根据你的实际路径修改以下变量 === INPUT_JSON = "/data/coding/coco/_annotations.coco.json" IMAGE_DIR = "/data/coding/coco" # 你图片所在的目录 OUTPUT_DIR = "/data/coding/coco_train" # 划分后的新根目录 split_coco_dataset(INPUT_JSON, IMAGE_DIR, OUTPUT_DIR, split_ratio=0.8)
开始训练：
from rfdetr import RFDETRBase, RFDETRNano, RFDETRSmall, RFDETRMedium # n/s/m/l # model = RFDETRNano() model = RFDETRSmall() # model = RFDETRMedium() # model = RFDETRBase() model.train( dataset_dir="/data/coding/coco_train", # 加载数据地址文件夹 epochs=200, batch_size=8, grad_accum_steps=4, lr=1e-4, resolution=512, output_dir="rf-detr_weight", # 模型训练完成的输出地址 # resume = "", early_stopping=True, tensorboard=True, pin_memory=True, num_workers=8 )
导出onnx:
from rfdetr import RFDETRSmall # 1. 实例化你训练时使用的模型缩放 # 2. pretrain_weights 指向你训练得到的 .pth 文件 model = RFDETRSmall(pretrain_weights="checkpoint_best_regular.pth") # 3. 执行导出 (默认会生成在当前目录下) model.export() print("ONNX 导出完成！")
6. 推理代码与效果
import time import cv2 import numpy as np import onnxruntime as ort def letterbox(im, new_shape=(512, 512), color=(114, 114, 114)): """保持比例缩放并进行补白 (Letterbox)，防止图像变形导致坐标偏移""" shape = im.shape[:2] # 当前形状 [height, width] if isinstance(new_shape, int): new_shape = (new_shape, new_shape) # 计算缩放比例 r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) # 计算缩放后的未填充形状 new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r))) dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # 填充的宽高 dw /= 2 # 两侧平分 dh /= 2 if shape[::-1] != new_unpad: # 调整大小 im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # 填充边框 return im, r, (left, top) def preprocess_image(image, input_size=512): """预处理：Letterbox -> 归一化 -> CHW""" if image is None: raise ValueError("输入图片为空") # 1. Letterbox 缩放 (保持比例) img_padded, ratio, (pad_w, pad_h) = letterbox(image, new_shape=input_size) # 2. 归一化 img_data = img_padded.astype(np.float32) / 255.0 # 3. HWC 转 CHW img_data = np.transpose(img_data, (2, 0, 1)) img_data = np.expand_dims(img_data, axis=0) return img_data, image.shape[:2], ratio, (pad_w, pad_h) def box_cxcywh_to_xyxy(x): """将 [center_x, center_y, width, height] 转换为 [x1, y1, x2, y2]""" cx, cy, w, h = x x1 = cx - 0.5 * w y1 = cy - 0.5 * h x2 = cx + 0.5 * w y2 = cy + 0.5 * h return np.array([x1, y1, x2, y2]) def inference_image(image, model_path="./output/inference_model.onnx", input_size=512, conf_threshold=0.5): """ 对单张图片进行推理 Args: image: numpy数组格式的图片 (BGR) model_path: ONNX模型路径 input_size: 模型输入尺寸 conf_threshold: 置信度阈值 Returns: detections: 检测结果列表，每个元素包含 [x1, y1, x2, y2, score, label] visualized_image: 绘制了检测框的图片 """ # --- 初始化 ONNX 会话 --- session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) input_name = session.get_inputs()[0].name # --- 预处理 --- img_data, (orig_h, orig_w), ratio, (pad_w, pad_h) = preprocess_image(image, input_size) # --- 推理 --- t1 = time.time() outputs = session.run(None, {input_name: img_data}) inference_time = time.time() - t1 # 识别输出节点 (2个节点的情况) if outputs[0].shape[-1] == 4: out_boxes, out_logits = outputs[0][0], outputs[1][0] else: out_logits, out_boxes = outputs[0][0], outputs[1][0] # --- 解析结果 --- detections = [] visualized_image = image.copy() for i in range(len(out_boxes)): probs = out_logits[i] label = np.argmax(probs) score = probs[label] if score > conf_threshold: # 1. 坐标格式转换 [cx,cy,w,h] -> [x1,y1,x2,y2] box = box_cxcywh_to_xyxy(out_boxes[i]) # 2. 还原 Letterbox 偏移和缩放 # 模型输出的是相对于 input_size 补齐图的归一化坐标 (0~1) x1 = (box[0] * input_size - pad_w) / ratio y1 = (box[1] * input_size - pad_h) / ratio x2 = (box[2] * input_size - pad_w) / ratio y2 = (box[3] * input_size - pad_h) / ratio # 3. 边界检查和四舍五入 x1, y1, x2, y2 = max(0, int(x1)), max(0, int(y1)), min(orig_w, int(x2)), min(orig_h, int(y2)) # 保存检测结果 detections.append([x1, y1, x2, y2, float(score), int(label)]) # 绘制检测框 cv2.rectangle(visualized_image, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.putText(visualized_image, f"Class:{label} {score:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) # 打印推理信息 print(f"推理耗时: {inference_time:.4f}s") print(f"检测到 {len(detections)} 个目标") for i, det in enumerate(detections): print(f" 目标 {i + 1}: 位置=({det[0]},{det[1]},{det[2]},{det[3]}), 置信度={det[4]:.3f}, 类别={det[5]}") return detections, visualized_image if __name__ == "__main__": # --- 配置 --- model_path = "./output/inference_model.onnx" image_path = r"./crop/test/images/capture_20250823_173410_896_original_840_840_1800_1800.jpg" input_size = 512 conf_threshold = 0.5 # --- 读取图片 --- image = cv2.imread(image_path) if image is None: print(f"错误: 无法读取图片 {image_path}") exit(1) # --- 推理 --- detections, result_image = inference_image( image, model_path=model_path, input_size=input_size, conf_threshold=conf_threshold ) # --- 保存结果 --- cv2.imwrite("./onnx_test_image4.png", result_image) print(f"结果已保存至: ./onnx_test_image4.png")
小结:借助AI,简单介绍了RF-DETR，并测试了自己的数据集，效果还不错。
浙公网安备 33010602011771号