RF-DETR实现onnx推理(Python版)
本文不生产技术，只做技术的搬运工！！！
前言
上篇文章完成了RF-DETR本地推理和自建数据集训练的介绍，但是推理太慢了，需要加一下速，本文章介绍如何实现onnx推理。训练及环境配置参考：
环境补充
配置训练环境时，有一些包没有安装，这里进行一些补充，就记得这几个了，其余的忘了，大家缺啥补啥吧，反正都是pip安装。
conda activate rf-detr pip install onnx pip install onnxruntime pip install onnxsimonnx模型导出
from rfdetr import RFDETRBase,RFDETRLarge,RFDETRNano,RFDETRSmall,RFDETRMedium model = RFDETRNano(pretrain_weights="/home/project_python/rf-detr/runs/nano/checkpoint_best_total.pth") model.export(output_dir = "/home/project_python/rf-detr/runs/nano/")onnx模型推理
这里作者实现了两个版本的预处理，一个使用torch一个使用numpy，经过测试，torch预处理的图像推理效果与pt模型推理效果一致，numpy版本推理结果与pt模型略有出入，但影响不大，由于工程端一般不使用torch，大家可以选择numpy版本的预处理，将其转换成C++代码进行部署。
import cv2 import onnxruntime import numpy as np import random import time import torchvision.transforms.functional as F def xywh2xyxy(x): """ 将边界框坐标从 (x, y, width, height) 转换为 (x1, y1, x2, y2) Args: x (np.ndarray): 输入边界框数组 Returns: np.ndarray: 转换后的边界框数组 """ y = np.copy(x) y[..., 0] = x[..., 0] - x[..., 2] / 2 # top left x y[..., 1] = x[..., 1] - x[..., 3] / 2 # top left y y[..., 2] = x[..., 0] + x[..., 2] / 2 # bottom right x y[..., 3] = x[..., 1] + x[..., 3] / 2 # bottom right y return y def softmax(x): """ Softmax函数 Args: x (np.ndarray): 输入数组 Returns: np.ndarray: 经过softmax处理的数组 """ e_x = np.exp(x - np.max(x)) return e_x / e_x.sum(axis=0) def sigmoid(x): """ Sigmoid函数 Args: x (np.ndarray): 输入数组 Returns: np.ndarray: 经过sigmoid处理的数组 """ return 1 / (1 + np.exp(-x)) def get_optimal_font_scale(image_shape, text, font_face=cv2.FONT_HERSHEY_SIMPLEX): """ 根据图像分辨率自动计算最优字体大小 """ # 获取图像的最小边长 min_dimension = min(image_shape[0], image_shape[1]) # 基于图像尺寸计算基础字体大小 base_font_scale = min_dimension / 1000.0 # 确保字体大小在合理范围内 font_scale = max(0.5, min(base_font_scale, 2.0)) return font_scale def generate_distinct_color(): """ 生成具有明显区分度的随机颜色 """ # 使用HSV色彩空间生成颜色，确保颜色具有较高的饱和度和亮度 h = random.randint(0, 179) # OpenCV中色调范围是0-179 s = random.randint(150, 255) # 饱和度：150-255 (避免过淡的颜色) v = random.randint(150, 255) # 亮度：150-255 (避免过暗的颜色) # 将HSV转换为BGR hsv = np.uint8([[[h, s, v]]]) bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) return tuple(int(x) for x in bgr[0][0]) def get_color_for_class(class_id, color_map): """ 为类别ID获取颜色，如果不存在则生成新颜色 """ if class_id not in color_map: color_map[class_id] = generate_distinct_color() return color_map[class_id] def preprocess_image(image, target_size=(384, 384)): """ 预处理图像 Args: image (np.ndarray): 输入图像 target_size (tuple): 目标图像尺寸 Returns: tuple: (处理后的图像, 缩放因子, 原始尺寸) """ # 保存原始尺寸 h, w = image.shape[:2] w_rate = w / target_size[0] h_rate = h / target_size[1] # 调整图像大小 resized_image = cv2.resize(image, target_size,interpolation=cv2.INTER_LINEAR) # 颜色空间转换和归一化 rgb_image = cv2.cvtColor(resized_image, cv2.COLOR_BGR2RGB) normalized_image = rgb_image.astype(np.float32) / 255.0 # 标准化 mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] normalized_image = (normalized_image - mean) / std # 转换维度并扩展批次维度 processed_image = np.transpose(normalized_image, (2, 0, 1)) processed_image = np.expand_dims(processed_image, axis=0).astype(np.float32) return processed_image, (w_rate, h_rate), (w, h) def preprocess_image_torch(image, target_size=(384, 384)): """ 预处理图像 Args: image (np.ndarray): 输入图像 target_size (tuple): 目标图像尺寸 Returns: tuple: (处理后的图像, 缩放因子, 原始尺寸) """ img_tensor = image img_tensor = F.to_tensor(img_tensor).to("cuda:0") h, w = img_tensor.shape[1:] mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] img_tensor = F.normalize(img_tensor, mean, std) img_tensor = F.resize(img_tensor, target_size) # 保存原始尺寸 h, w = image.shape[:2] w_rate = w / target_size[0] h_rate = h / target_size[1] return np.expand_dims(img_tensor.cpu().numpy(), axis=0).astype(np.float32), (w_rate, h_rate), (w, h) def postprocess_detections(bboxes, labels, target_size=(384, 384)): """ 后处理检测结果 Args: bboxes (np.ndarray): 边界框数组 labels (np.ndarray): 标签数组 target_size (tuple): 目标图像尺寸 Returns: list: 处理后的检测结果列表 """ results = [] for i in range(len(bboxes)): # 转换边界框格式 bbox = xywh2xyxy(bboxes[i]) # 缩放边界框坐标 bbox[0] *= target_size[1] # x坐标乘以宽度 bbox[2] *= target_size[1] # x坐标乘以宽度 bbox[1] *= target_size[0] # y坐标乘以高度 bbox[3] *= target_size[0] # y坐标乘以高度 # 应用softmax并获取类别和置信度 label_list = sigmoid(labels[i]) class_id = np.argmax(label_list) conf = label_list[class_id] results.append({ 'bbox': bbox, 'class_id': class_id, 'confidence': conf }) return results def draw_detections(image, detections, scale_factors, conf_threshold=0.5): """ 在图像上绘制检测结果 Args: image (np.ndarray): 输入图像 detections (list): 检测结果列表 scale_factors (tuple): 缩放因子 (w_rate, h_rate) conf_threshold (float): 置信度阈值 Returns: np.ndarray: 绘制了检测结果的图像 """ w_rate, h_rate = scale_factors result_image = image.copy() detection_count = 0 # 存储类别颜色的字典 color_map = {} # 根据图像分辨率自动调整字体大小 font_scale = get_optimal_font_scale(image.shape, "SampleText") thickness = max(1, int(font_scale * 2)) # 根据字体大小调整线条粗细 for detection in detections: conf = detection['confidence'] if conf > conf_threshold: detection_count += 1 bbox = detection['bbox'] class_id = detection['class_id'] # 为每个类别ID获取颜色 color = get_color_for_class(class_id, color_map) # 在原始尺寸图像上绘制框 x1 = int(bbox[0] * w_rate) y1 = int(bbox[1] * h_rate) x2 = int(bbox[2] * w_rate) y2 = int(bbox[3] * h_rate) # 绘制边界框 cv2.rectangle(result_image, (x1, y1), (x2, y2), color, thickness) # 在同一行显示类别名称和置信度，带背景色 label = f"{class_id} {conf:.2f}" # 计算文本尺寸 (text_width, text_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness) # 绘制文本背景框 cv2.rectangle(result_image, (x1, y1 - text_height - baseline - 2), (x1 + text_width, y1), color, -1) # 绘制文本 cv2.putText(result_image, label, (x1, y1 - baseline), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thickness) print(f"检测框: 类别={class_id}, 置信度={conf:.2f}") print(f"总共检测到 {detection_count} 个置信度高于阈值的目标") return result_image def run_detection(model_path, image_path, conf_threshold=0.5): """ 运行目标检测 Args: model_path (str): ONNX模型路径 image_path (str): 图像路径 conf_threshold (float): 置信度阈值 """ # 初始化模型 model = onnxruntime.InferenceSession(model_path,providers=['CPUExecutionProvider']) # 读取图像 image = cv2.imread(image_path) src_image = image.copy() # 预处理图像 processed_image, scale_factors, original_size = preprocess_image(image) # 模型推理 start = time.time() output = model.run(["dets", "labels"], {"input": processed_image}) print(f"inference time: {time.time() - start}") bboxes = output[0][0] labels = output[1][0] # 后处理检测结果 detections = postprocess_detections(bboxes, labels) # 绘制检测结果 result_image = draw_detections(src_image, detections, scale_factors, conf_threshold) # 保存结果 cv2.imwrite("result_onnx.jpg", result_image) # 主程序 if __name__ == "__main__": model_path = "/home/project_python/rf-detr/runs/nano/inference_model.onnx" image_path = "/home/project_python/rf-detr/dataset/dtrain_20251030_v1/images/test/gacdbx_1_17424bef_1600_20251029113538995.jpg" start = time.time() run_detection(model_path, image_path, conf_threshold=0.9) print(f"total time: {time.time() - start}")结语
onnx模型推理速度为72ms每张，下篇文章将更新tensorrt推理。
