DETR模型ONNX推理实战：从输出张量到可视化检测框的完整解析

DETR模型ONNX推理实战从输出张量到可视化检测框的完整解析在目标检测领域DETRDetection Transformer以其独特的端到端架构和简洁的流程设计正在改变传统基于锚框anchor-based方法的格局。不同于Faster R-CNN等经典检测模型DETR直接输出固定数量的预测结果默认为100个省去了繁琐的锚框设计和后处理步骤。然而这种创新也带来了新的挑战——如何正确解析模型输出的pred_logits和pred_boxes张量并将其转换为直观的可视化检测结果这正是本文要解决的核心问题。我们将从ONNX推理的实际场景出发深入剖析DETR的输出结构手把手演示如何将归一化的中心坐标宽高cxcywh转换为图像上的绝对坐标xyxy并实现置信度过滤与结果可视化。无论您是在部署模型时遇到输出解析困惑还是希望深入理解DETR的工作原理本文都将为您提供清晰的解决方案和可复用的代码工具。1. DETR输出结构深度解析DETR模型的输出包含两个关键张量pred_logits和pred_boxes。理解它们的结构和含义是正确解析结果的第一步。1.1 输出张量的维度与含义典型的DETR ONNX模型输出如下pred_logits形状为[1, 100, 92]的张量1批处理大小batch size100固定数量的预测框92COCO数据集的类别数91个物体类别1个背景类pred_boxes形状为[1, 100, 4]的张量4每个预测框的坐标参数格式为[中心x, 中心y, 宽度, 高度]即cxcywh所有坐标值都是相对于输入图像尺寸的归一化值范围0-1# 示例查看输出张量的形状 import onnxruntime as rt sess rt.InferenceSession(detr.onnx) output_names [output.name for output in sess.get_outputs()] net_outs sess.run(output_names, {images: input_tensor}) print(pred_logits shape:, net_outs[0].shape) # (1, 100, 92) print(pred_boxes shape:, net_outs[1].shape) # (1, 100, 4)1.2 预测框的分布特点DETR的100个预测框中大部分是低质量的空预测对应背景类。实际应用中我们需要通过置信度过滤掉低质量预测对剩余预测框进行非极大值抑制NMS去除冗余下表展示了典型DETR输出的统计特性预测框类型占比典型置信度范围处理方式高质量检测5-15%0.7保留并可视化中等质量检测10-20%0.3-0.7根据应用场景选择保留低质量/背景65-85%0.3过滤掉2. 坐标转换从cxcywh到图像像素坐标DETR输出的pred_boxes采用[中心x, 中心y, 宽度, 高度]的归一化格式需要转换为图像上的绝对坐标[x_min, y_min, x_max, y_max]才能用于可视化。2.1 归一化坐标到绝对坐标的转换转换过程分为两步将cxcywh转换为xyxy格式仍为归一化坐标根据图像尺寸将归一化坐标缩放为像素坐标import numpy as np def box_cxcywh_to_xyxy(x): 将cxcywh格式的框转换为xyxy格式 x_c, y_c, w, h x[..., 0], x[..., 1], x[..., 2], x[..., 3] b [(x_c - 0.5 * w), # x_min (y_c - 0.5 * h), # y_min (x_c 0.5 * w), # x_max (y_c 0.5 * h)] # y_max return np.stack(b, axis-1) def scale_boxes(boxes, image_size): 将归一化坐标缩放为图像像素坐标 height, width image_size scale_fct np.array([width, height, width, height]) return boxes * scale_fct2.2 处理不同尺寸的输入图像实际应用中输入图像可能经过resize处理。为确保坐标转换正确需要记录原始图像尺寸和预处理方式# 示例处理经过resize的图像 original_image Image.open(input.jpg) original_size original_image.size # (width, height) # 预处理resize到512x512 processed_image original_image.resize((512, 512)) input_tensor preprocess(processed_image) # 假设已实现预处理函数 # 模型推理后转换坐标时使用原始尺寸 boxes box_cxcywh_to_xyxy(pred_boxes) scaled_boxes scale_boxes(boxes, original_size)3. 类别分数处理与置信度过滤pred_logits包含每个预测框的类别分数但需要经过适当处理才能得到可用的置信度分数。3.1 从logits到概率分布DETR输出的logits需要通过softmax转换为概率分布def process_logits(pred_logits): # 对最后一个维度类别维度应用softmax prob np.exp(pred_logits) / np.exp(pred_logits).sum(-1, keepdimsTrue) # 获取每个框的最高分数和对应类别忽略背景类 scores np.max(prob[..., :-1], axis-1) # 去掉背景类 labels np.argmax(prob[..., :-1], axis-1) return scores, labels3.2 置信度阈值的选择策略选择合适的置信度阈值对结果质量至关重要高阈值0.7以上只保留高质量预测减少假阳性中等阈值0.3-0.7平衡召回率和准确率低阈值0.3以下适合需要高召回率的场景提示对于实时应用建议从0.5开始调整根据实际效果微调阈值。4. 完整推理流程与可视化实现将上述步骤整合为一个完整的推理流程并实现检测结果的可视化。4.1 构建解码工具类class DETRDecoder: def __init__(self, class_names): self.class_names class_names def decode_output(self, pred_logits, pred_boxes, image_size, confidence_thresh0.5): # 处理类别分数 prob np.exp(pred_logits) / np.exp(pred_logits).sum(-1, keepdimsTrue) scores np.max(prob[..., :-1], axis-1) labels np.argmax(prob[..., :-1], axis-1) # 转换框格式 boxes self.box_cxcywh_to_xyxy(pred_boxes) boxes self.scale_boxes(boxes, image_size) # 过滤低置信度预测 keep scores[0] confidence_thresh filtered_boxes boxes[0][keep] filtered_scores scores[0][keep] filtered_labels labels[0][keep] return filtered_boxes, filtered_scores, filtered_labels def box_cxcywh_to_xyxy(self, x): x_c, y_c, w, h x[..., 0], x[..., 1], x[..., 2], x[..., 3] b [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c 0.5 * w), (y_c 0.5 * h)] return np.stack(b, axis-1) def scale_boxes(self, boxes, image_size): height, width image_size[1], image_size[0] scale_fct np.array([width, height, width, height]) return boxes * scale_fct4.2 可视化检测结果from PIL import ImageDraw, ImageFont import colorsys def visualize_detections(image, boxes, scores, labels, class_names): draw ImageDraw.Draw(image) font ImageFont.load_default() # 为每个类别生成颜色 num_classes len(class_names) hsv_tuples [(x / num_classes, 1., 1.) for x in range(num_classes)] colors list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples)) colors list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors)) for box, score, label in zip(boxes, scores, labels): # 绘制边界框 draw.rectangle(box.tolist(), outlinecolors[label], width2) # 准备标签文本 label_text f{class_names[label]}: {score:.2f} label_size draw.textsize(label_text, font) # 绘制标签背景 text_origin (box[0], box[1] - label_size[1]) draw.rectangle( [text_origin[0], text_origin[1], text_origin[0] label_size[0], text_origin[1] label_size[1]], fillcolors[label] ) # 绘制文本 draw.text(text_origin, label_text, fill(0, 0, 0), fontfont) return image4.3 端到端推理示例# 加载模型和类名 sess rt.InferenceSession(detr.onnx) class_names [...] # 加载COCO类名 # 初始化解码器 decoder DETRDecoder(class_names) # 预处理输入图像 image Image.open(input.jpg) original_size image.size input_tensor preprocess_image(image) # 实现预处理函数 # 运行推理 outputs sess.run(None, {images: input_tensor}) pred_logits, pred_boxes outputs[0], outputs[1] # 解码输出 boxes, scores, labels decoder.decode_output( pred_logits, pred_boxes, original_size, confidence_thresh0.5 ) # 可视化并保存结果 result_image visualize_detections(image, boxes, scores, labels, class_names) result_image.save(output.jpg)在实际项目中这种模块化的设计使得DETR模型的集成变得简单高效。解码器类可以轻松集成到各种应用场景中而可视化函数则可以根据具体需求进行定制比如添加不同的框样式、调整标签显示方式等。