1. DETR模型预测结果自动化导出的核心价值第一次用DETR做目标检测时最让我头疼的就是如何把预测结果整理成结构化数据。每次跑完模型面对满屏的检测框坐标和类别信息手动记录简直是一场噩梦。直到摸索出这套自动化导出方案工作效率直接提升10倍不止。DETR作为Transformer架构在目标检测领域的代表作其端到端的特性让模型输出非常干净——每张图片的预测结果都包含三个关键要素物体类别、置信度分数、边界框坐标cxcywh格式。但官方代码默认只提供可视化功能要想批量处理数百张图片并生成标注文件就需要对原始代码进行外科手术式改造。这个自动化流程特别适合需要批量评估模型性能的算法工程师准备将检测结果导入其他系统的开发人员制作新数据集的标注人员我曾用这套方法在智慧工地安全帽检测项目中3小时就完成了2000张测试图片的结果导出生成的TXT文件直接对接后续的MES系统。下面就把完整实现方案拆解给你看。2. 环境准备与代码改造2.1 基础环境搭建建议使用Python 3.8和PyTorch 1.10的组合这是经过实测最稳定的版本搭配。如果要用GPU加速别忘了安装对应版本的CUDA工具包。我的测试环境配置如下conda create -n detr python3.8 conda install pytorch1.12.1 torchvision0.13.1 cudatoolkit11.3 -c pytorch pip install matplotlib opencv-python2.2 核心代码改造要点原始DETR的demo代码只能单张图片可视化我们要改造三个关键部分批量图片处理用os.listdir遍历整个图片文件夹结果文件存储创建独立的labels目录存放TXT标注文件坐标格式转换将cxcywh转为实际像素坐标这里有个坑要注意DETR输出的边界框坐标是归一化的cxcywh格式中心点坐标宽高而常见的标注工具需要的是实际像素坐标。改造后的关键函数如下def rescale_bboxes(out_bbox, width, height): box_coords box_cxcywh_to_xyxy(out_bbox) scale_tensor torch.Tensor([width, height, width, height]).to(device) return box_coords * scale_tensor3. 完整实现流程3.1 文件目录结构设计规范的目录结构能让后续处理事半功倍。建议按这个方式组织project_root/ ├── inputs/ # 存放待检测图片 ├── outputs/ │ ├── images/ # 保存带检测框的可视化图片 │ └── labels/ # 保存TXT标注文件 └── predict.py # 改造后的预测脚本3.2 核心代码实现以下是改造后的完整预测代码框架import os from PIL import Image import torch import torchvision.transforms as T from util.misc import nested_tensor_from_tensor_list # 初始化模型 model detr_resnet50(pretrainedFalse, num_classes91) checkpoint torch.load(checkpoint.pth, map_locationcuda) model.load_state_dict(checkpoint[model]) model.to(device) model.eval() # 图片预处理流水线 transform T.Compose([ T.Resize(800), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) def save_results(txt_path, class_id, confidence, bbox): 将检测结果写入TXT文件 with open(txt_path, a) as f: line f{class_id} {confidence} {bbox[0]} {bbox[1]} {bbox[2]} {bbox[3]}\n f.write(line) # 批量处理图片 for img_name in os.listdir(inputs): img_path os.path.join(inputs, img_name) img Image.open(img_path) # 执行预测 outputs model(transform(img).unsqueeze(0).to(device)) # 处理预测结果 probas outputs[pred_logits].softmax(-1)[0, :, :-1] keep probas.max(-1).values 0.5 # 置信度阈值 # 保存结果 txt_path os.path.join(outputs/labels, f{os.path.splitext(img_name)[0]}.txt) for p, box in zip(probas[keep], outputs[pred_boxes][0, keep]): class_id p.argmax() confidence p[class_id].item() bbox rescale_bboxes(box.unsqueeze(0), img.width, img.height)[0].tolist() save_results(txt_path, class_id, confidence, bbox)4. 实战技巧与避坑指南4.1 置信度阈值调优DETR的预测结果包含大量低置信度的检测框需要合理设置阈值过滤。根据我的经验室内场景建议0.5-0.6复杂室外场景建议0.3-0.4对小物体检测可降低到0.2可以通过这个代码动态调整keep probas.max(-1).values threshold # 调整threshold值4.2 内存泄漏问题处理当处理大量图片时plt画图可能导致内存泄漏。我的解决方案是强制回收figure资源使用agg后端避免交互式显示import matplotlib matplotlib.use(agg) # 在import pyplot前设置 plt.close(all) # 每次画图后调用4.3 多GPU环境适配如果在多GPU服务器上训练模型加载checkpoint时需要特殊处理# 原始状态字典的键名可能包含module.前缀 state_dict checkpoint[model] new_state_dict {k.replace(module., ): v for k,v in state_dict.items()} model.load_state_dict(new_state_dict)5. 结果验证与应用5.1 标注文件格式解析生成的TXT文件每行对应一个检测物体格式为类别ID 置信度 x_min y_min x_max y_max例如2 0.98 354 120 402 158 3 0.87 102 205 145 2405.2 与其他工具集成这种标准化格式可以轻松转换为其他格式转COCO格式用json.dump重写标注转YOLO格式将坐标归一化为相对值导入LabelImg通过脚本批量转换这里提供一个转YOLO格式的示例def convert_to_yolo(img_w, img_h, bbox): x_center (bbox[0] bbox[2]) / 2 / img_w y_center (bbox[1] bbox[3]) / 2 / img_h width (bbox[2] - bbox[0]) / img_w height (bbox[3] - bbox[1]) / img_h return [x_center, y_center, width, height]在实际工业质检项目中这套自动化流程将原本需要3天的手动标注工作压缩到2小时内完成。特别是在处理玻璃缺陷检测这类小物体场景时通过调整DETR的encoder层数配合这套导出方案使mAP指标提升了15个百分点。