1. 为什么选择DETR从零开始的理由如果你之前用过YOLO或者Faster R-CNN这类目标检测模型第一次接触DETR时可能会有点懵。它没有那些复杂的锚框Anchor Boxes也不需要非极大值抑制NMS这种后处理步骤整个模型结构看起来干净利落。我刚开始用的时候也觉得这玩意儿是不是太“理想化”了但实际在几个自定义项目里跑下来尤其是在处理类别多、目标大小差异大的场景时DETR的表现真的让我改观了。DETR的核心思想是把目标检测当成一个“集合预测”问题。你可以把它想象成一场考试Transformer编码器-解码器结构就是那个考生它“看”完整张图片后直接输出固定数量比如100个的预测结果每个结果包含一个类别和一个边界框。这100个预测里有些对应真实物体有些则对应“无物体”。模型通过一个独特的“二分图匹配”损失函数在学习过程中自己学会把预测和真实目标一一对应起来。这种端到端的方式最大的好处就是流程简化了你不需要再去调那些锚框的尺寸、宽高比也不用担心NMS阈值设高了漏检、设低了重复框的问题。那么什么时候你应该考虑用DETR来训练自己的数据集呢从我踩过的坑来看主要有这么几种情况一是你的数据集中目标类别特别多像我做过的那个路况数据集有57类用传统方法设计锚框会非常头疼二是你的目标形状、尺度变化很大DETR的全局注意力机制在这方面有天然优势三是你追求更干净、更现代的 pipeline不想在数据预处理和后处理上花太多心思。当然DETR也不是万能的它对计算资源要求比较高训练时间通常比YOLO系列要长而且在小目标检测上早期的DETR效果可能不如一些精心设计的单阶段检测器。不过随着后续改进版本如Deformable DETR的出现很多问题都已经得到了缓解。2. 实战第一步理解你的数据与COCO格式要训练DETR第一步不是急着写代码而是要把你的数据“收拾”明白。DETR官方代码默认支持的是COCO数据集格式所以我们的目标就是把手里五花八门的数据都转换成COCO的“样子”。这就像你要把不同国家的电器插头都转换成一个标准的接口才能插上电源。COCO格式的核心是一个JSON文件它就像一个结构清晰的目录册。这个JSON文件里主要包含三大部分images、annotations和categories。images数组记录每张图片的信息比如文件名、宽度、高度和一个唯一的ID。annotations数组是重头戏它保存了所有标注框的信息每个标注框都要关联到对应的图片ID和类别ID并且要提供边界框坐标注意COCO用的是[x_min, y_min, width, height]格式而且是绝对坐标不是相对坐标。categories数组则定义了数据集中所有类别的名称和ID。我遇到最常见的数据格式是YOLO格式的TXT文件每行是class_id x_center y_center width height这里的坐标都是相对于图片宽度和高度的归一化值0到1之间。另一个常见的是VOC格式的XML文件。从YOLO格式转换到COCO格式关键的一步就是坐标转换。你需要把归一化的中心点坐标和宽高还原成图片上的绝对像素坐标。这个计算不复杂但一定要细心我早期就因为在转换时宽高顺序弄反了导致训练时损失一直不下降排查了好久。这里有个特别容易出错的地方类别ID的连续性。COCO格式要求类别ID通常从1开始连续编号0有时预留给背景。而YOLO的TXT里class_id很可能也是从0开始的。你在转换时一定要建立一个明确的映射关系并确保在最终的JSON文件里categories数组的顺序和ID与每个标注框里的category_id能完美对应上。我建议在转换脚本的开头就先把你的类别名称列表固定下来比如用一个Python列表然后基于这个列表来生成ID这样最不容易乱。3. 手把手构建从YOLO格式到COCO格式的转换脚本光讲原理不够我们直接来看一个我修改过、在实际项目中用烂了的转换脚本。这个脚本能帮你把YOLO格式的标注文件夹变成一个标准的COCO格式JSON文件。我会把关键步骤拆开加上详细的注释你完全可以拿来即用只需要改改文件路径。首先你需要准备好你的数据文件夹。假设你的数据是这样组织的your_dataset/ ├── images/ │ ├── train/ │ │ ├── 000001.jpg │ │ └── ... │ └── val/ │ ├── 000100.jpg │ └── ... └── labels/ ├── train/ │ ├── 000001.txt │ ├── classes.txt │ └── ... └── val/ ├── 000100.txt └── classes.txt注意在每个labels/train和labels/val文件夹下都必须放一个classes.txt文件这个文件按行写入你所有的类别名称。这是脚本读取类别列表的依据。下面是完整的Python转换脚本我把它保存为yolo_to_coco.pyimport argparse import os import sys import json import shutil import cv2 from datetime import datetime # 初始化一个空的COCO格式字典 coco dict() coco[images] [] coco[type] instances coco[annotations] [] coco[categories] [] category_set dict() image_set set() # 用于记录已添加的图片避免重复 image_id 0 # 图片ID计数器从1开始递增 annotation_id 0 # 标注ID计数器从1开始递增 def addCatItem(category_dict): 将类别字典添加到coco[categories]中 for k, v in category_dict.items(): category_item dict() category_item[supercategory] none # 通常设为none category_item[id] int(k) # 类别ID注意转换为整数 category_item[name] v # 类别名称 coco[categories].append(category_item) def addImgItem(file_name, size): 添加图片信息到coco[images]中 global image_id image_id 1 image_item dict() image_item[id] image_id image_item[file_name] file_name image_item[width] size[1] # 注意OpenCV读取的shape是(height, width, channels) image_item[height] size[0] image_item[license] None image_item[flickr_url] None image_item[coco_url] None image_item[date_captured] str(datetime.today()) coco[images].append(image_item) image_set.add(file_name) return image_id # 返回当前图片的ID用于关联标注 def addAnnoItem(object_name, image_id, category_id, bbox): 添加一个标注框信息到coco[annotations]中 global annotation_id annotation_item dict() annotation_item[segmentation] [] # 实例分割信息目标检测可留空列表 # 这里为了格式完整根据bbox生成了一个简单的矩形segmentation seg [] # bbox格式是 [x_min, y_min, width, height] x, y, w, h bbox seg.append(x) seg.append(y) seg.append(x) seg.append(y h) seg.append(x w) seg.append(y h) seg.append(x w) seg.append(y) annotation_item[segmentation].append(seg) annotation_item[area] w * h # 边界框面积 annotation_item[iscrowd] 0 # 是否为拥挤群组默认为0 annotation_item[ignore] 0 annotation_item[image_id] image_id annotation_item[bbox] bbox # [x_min, y_min, width, height] annotation_item[category_id] category_id annotation_id 1 annotation_item[id] annotation_id coco[annotations].append(annotation_item) def xywhn2xywh(bbox, img_size): 将YOLO格式的归一化坐标 [x_center, y_center, width, height] 转换为COCO格式的绝对坐标 [x_min, y_min, width, height] bbox: 包含四个字符串的列表如 [0.5, 0.5, 0.2, 0.3] img_size: 图片的尺寸格式为 (height, width) bbox list(map(float, bbox)) h, w img_size # 注意顺序 # 计算绝对坐标 x_center_abs bbox[0] * w y_center_abs bbox[1] * h width_abs bbox[2] * w height_abs bbox[3] * h # 计算左上角坐标 x_min x_center_abs - (width_abs / 2.0) y_min y_center_abs - (height_abs / 2.0) # 转换为整数并确保不超出图像边界 x_min max(0, int(round(x_min))) y_min max(0, int(round(y_min))) width_abs int(round(width_abs)) height_abs int(round(height_abs)) # 确保边界框不超出图像右边界和下边界 if x_min width_abs w: width_abs w - x_min if y_min height_abs h: height_abs h - y_min return [x_min, y_min, width_abs, height_abs] def parseTxtFiles(image_dir, label_dir, output_dir, json_name): 主转换函数 # 检查路径是否存在 assert os.path.exists(image_dir), f图片目录不存在: {image_dir} assert os.path.exists(label_dir), f标注目录不存在: {label_dir} # 创建输出目录 if os.path.exists(output_dir): shutil.rmtree(output_dir) os.makedirs(output_dir) json_path os.path.join(output_dir, json_name) # 1. 读取类别列表 class_file os.path.join(label_dir, classes.txt) assert os.path.exists(class_file), f类别文件不存在: {class_file} category_list [] with open(class_file, r) as f: for line in f.readlines(): category_list.append(line.strip()) # 构建类别ID到名称的映射通常ID从1开始0留给背景 category_id_dict {i1: name for i, name in enumerate(category_list)} addCatItem(category_id_dict) print(f成功加载 {len(category_list)} 个类别。) # 2. 获取所有图片和标签文件 image_files [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.lower().endswith((.jpg, .jpeg, .png))] label_files [os.path.join(label_dir, f) for f in os.listdir(label_dir) if f.endswith(.txt) and f ! classes.txt] # 建立文件名不含后缀到图片路径的映射方便查找 image_index {} for img_path in image_files: file_name os.path.basename(img_path) base_name os.path.splitext(file_name)[0] image_index[base_name] img_path # 3. 遍历每个标签文件 processed_count 0 for label_file in label_files: base_name os.path.splitext(os.path.basename(label_file))[0] # 查找对应的图片文件 if base_name not in image_index: print(f警告找不到标签文件 {label_file} 对应的图片已跳过。) continue img_path image_index[base_name] # 读取图片获取尺寸 img cv2.imread(img_path) if img is None: print(f警告无法读取图片 {img_path}已跳过。) continue h, w img.shape[:2] file_name os.path.basename(img_path) # 添加图片信息并获取其唯一ID current_image_id addImgItem(file_name, (h, w)) # 读取标签文件中的每一行每个目标 with open(label_file, r) as f: lines f.readlines() for line in lines: line line.strip() if not line: continue parts line.split() if len(parts) ! 5: print(f警告标签文件 {label_file} 中存在格式错误的行: {line}已跳过。) continue class_id int(parts[0]) # YOLO格式的类别ID通常从0开始 # 转换为COCO格式的类别ID通常1 coco_category_id class_id 1 # 坐标转换 bbox_normalized parts[1:5] bbox_abs xywhn2xywh(bbox_normalized, (h, w)) # 获取类别名称用于调试 category_name category_id_dict.get(coco_category_id, funknown_{coco_category_id}) # 添加标注信息 addAnnoItem(category_name, current_image_id, coco_category_id, bbox_abs) processed_count 1 if processed_count % 100 0: print(f已处理 {processed_count} 个文件...) # 4. 将结果保存为JSON文件 with open(json_path, w) as f: json.dump(coco, f, indent2) # 使用indent让JSON文件更易读 print(*50) print(f转换完成) print(f类别数量: {len(coco[categories])}) print(f图片数量: {len(coco[images])}) print(f标注框数量: {len(coco[annotations])}) print(fJSON文件已保存至: {json_path}) print(*50) if __name__ __main__: parser argparse.ArgumentParser(description将YOLO格式标注转换为COCO格式JSON文件。) parser.add_argument(--image-dir, requiredTrue, help图片文件夹路径如 ./images/train) parser.add_argument(--label-dir, requiredTrue, helpYOLO标签文件夹路径必须包含classes.txt文件如 ./labels/train) parser.add_argument(--output-dir, default./annotations, helpJSON文件输出目录) parser.add_argument(--json-name, defaultinstances_train2017.json, help输出的JSON文件名) args parser.parse_args() # 执行转换 parseTxtFiles(args.image_dir, args.label_dir, args.output_dir, args.json_name)使用这个脚本非常简单。假设你的训练集图片在/home/user/data/images/train标签在/home/user/data/labels/train你只需要在终端运行python yolo_to_coco.py \ --image-dir /home/user/data/images/train \ --label-dir /home/user/data/labels/train \ --output-dir /home/user/data/annotations \ --json-name instances_train2017.json对验证集或测试集再执行一次只需更改路径和输出的JSON文件名即可比如改成instances_val2017.json。运行成功后你的annotations文件夹里就会有两个标准的COCO格式JSON文件图片文件夹结构保持不变。这就是DETR需要的数据格式。4. 环境搭建与DETR代码准备数据准备好了接下来就得把DETR的训练环境搭起来。我强烈建议使用Anaconda来管理Python环境这样可以避免各种包版本冲突的“玄学”问题。DETR官方代码基于PyTorch所以我们需要先安装合适版本的PyTorch和Torchvision。首先创建一个新的conda环境我习惯叫它detrconda create -n detr python3.8 -y conda activate detr接着根据你的CUDA版本安装PyTorch。你可以去PyTorch官网查看最新的安装命令。假设你用的是CUDA 11.3可以这样安装pip install torch1.12.1cu113 torchvision0.13.1cu113 --extra-index-url https://download.pytorch.org/whl/cu113如果没有GPU就安装CPU版本。然后安装其他必要的依赖pip install cython scipy pip install opencv-python pillow matplotlib tqdm # 安装pycocotools用于COCO数据集评估 pip install pycocotools环境搞定后去DETR的官方GitHub仓库把代码克隆下来git clone https://github.com/facebookresearch/detr.git cd detr官方的代码结构比较清晰我们主要关心的文件是main.py训练和评估的主入口、models/detr.py模型定义以及datasets/coco.py数据加载。在开始训练前我们得先准备好预训练权重。DETR从零开始训练非常慢官方提供了在COCO数据集上预训练好的模型。我们可以用ResNet-50作为主干的预训练模型# 在detr目录下 wget https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth下载下来的detr-r50-e632da11.pth就是我们的预训练权重。但是这个权重文件最后的分类头是针对COCO的91类80个物体类1个背景设计的。我们的自定义数据集类别数完全不同所以不能直接使用必须对它进行修改。5. 关键改造适配自定义数据集的预训练权重与模型这是整个流程里最关键、也最容易出错的一步。我们需要修改预训练权重文件里的分类头参数让它匹配我们数据集的类别数。DETR的模型输出包含两个部分一个是分类分数class logits一个是边界框坐标bbox coordinates。分类头的权重矩阵形状是[num_classes1, hidden_dim]这里的1是给“无物体”no object这个特殊类别留的。假设我们的路况数据集有57个类别那么num_classes就是57加上“无物体”类总共需要58个输出。而下载的预训练权重里这个矩阵的形状是[92, 256]COCO的91类1。我们需要把这个矩阵的大小调整到[58, 256]。注意我们不能简单地截取或填充因为新矩阵的参数需要重新初始化。通常的做法是保留原来矩阵的前58行对应COCO的前57类和我们新增的“无物体”类但这只是一种启发式方法更好的做法可能是随机初始化新矩阵然后只加载能匹配上的部分。我写了一个简单的Python脚本modify_pretrain.py来做这件事import torch import sys def modify_pretrained_weights(pretrained_path, num_new_classes, output_path): 修改DETR预训练权重文件的分类头以适应新的类别数。 pretrained_path: 原始预训练权重路径 num_new_classes: 你的数据集的实际物体类别数不包括背景 output_path: 修改后权重的保存路径 # 加载预训练权重 print(f加载预训练权重: {pretrained_path}) pretrained_weights torch.load(pretrained_path, map_locationcpu) # 新的总类别数 物体类别数 1 (for no object class) new_num_classes num_new_classes 1 print(f新的类别数 (含no object): {new_num_classes}) # 获取原始分类头权重和偏置 class_embed_weight pretrained_weights[model][class_embed.weight] # 形状 [92, 256] class_embed_bias pretrained_weights[model][class_embed.bias] # 形状 [92] old_num_classes class_embed_weight.size(0) hidden_dim class_embed_weight.size(1) print(f原始类别数: {old_num_classes}, 隐藏层维度: {hidden_dim}) # 创建新的权重矩阵 new_class_weight torch.randn(new_num_classes, hidden_dim) * 0.01 # 小随机数初始化 new_class_bias torch.randn(new_num_classes) * 0.01 # 启发式策略如果新类别数小于等于旧类别数尝试保留前N行对应COCO常见类 # 注意这只是一种策略因为类别语义可能完全不同重新训练分类头是更稳妥的做法。 # 这里我们选择保留前 min(new_num_classes, old_num_classes) 行。 copy_len min(new_num_classes, old_num_classes) new_class_weight[:copy_len] class_embed_weight[:copy_len].detach() new_class_bias[:copy_len] class_embed_bias[:copy_len].detach() print(f复制了前 {copy_len} 个类别的权重和偏置。) # 替换模型字典中的权重 pretrained_weights[model][class_embed.weight] new_class_weight pretrained_weights[model][class_embed.bias] new_class_bias # 保存修改后的权重 torch.save(pretrained_weights, output_path) print(f修改后的权重已保存至: {output_path}) return output_path if __name__ __main__: # 参数设置 original_weights detr-r50-e632da11.pth # 你下载的预训练权重文件名 my_class_num 57 # 你的数据集的物体类别数例如路况数据集的57 output_weights fdetr_r50_{my_class_num1}.pth # 输出文件名 modify_pretrained_weights(original_weights, my_class_num, output_weights)运行这个脚本python modify_pretrain.py。它会生成一个新的权重文件比如detr_r50_58.pth。这个文件就是我们接下来训练要用的起点。别忘了还需要修改模型定义文件打开detr/models/detr.py找到定义DETR类的地方大约在300行附近你会看到num_classes这个参数。你需要把它改成你的类别数57。通常这行代码长这样num_classes91你把它改成num_classes57。这个修改确保了模型结构本身输出通道数是对的。6. 配置与启动开始训练你的DETR模型万事俱备只欠训练。DETR的训练通过main.py脚本启动它有一大堆参数可以配置。对于新手我们重点关注几个最关键的。首先确保你的数据集按照COCO格式组织好了目录结构类似这样coco_custom/ ├── annotations/ │ ├── instances_train2017.json │ └── instances_val2017.json └── images/ ├── train2017/ │ ├── 000001.jpg │ └── ... └── val2017/ ├── 000100.jpg └── ...注意images下的子文件夹名字train2017和val2017是DETR的coco.py里默认寻找的如果你用了别的名字需要在代码里修改路径映射或者直接创建符号链接。接下来在终端进入detr目录执行训练命令。一个比较基础的训练命令如下python main.py \ --dataset_file coco \ --coco_path /path/to/your/coco_custom \ --output_dir /path/to/save/checkpoints \ --resume detr_r50_58.pth \ --epochs 300 \ --lr 1e-4 \ --lr_backbone 1e-5 \ --batch_size 4 \ --num_workers 4 \ --device cuda \ --seed 42我来解释一下这些参数--dataset_file coco指定使用COCO格式的数据集加载器。--coco_path指向你自定义数据集的根目录即上面coco_custom的路径。--output_dir训练过程中模型检查点checkpoint和日志的保存目录。--resume我们刚才修改好的预训练权重路径。注意这里不是从某个训练轮次继续而是加载预训练权重作为初始化。--epochs训练总轮数。DETR通常需要训练较长时间300轮是官方常用的设置。--lr和--lr_backbone学习率。主干网络ResNet通常用更小的学习率lr_backbone因为它是预训练的微调即可。Transformer部分用较大的学习率。--batch_size根据你的GPU显存来调整。RTX 309024GB上可以尝试8或16显存小的话就调成2或4。--num_workers数据加载的线程数可以加快数据读取速度通常设为CPU核心数。--device指定使用GPUcuda还是CPUcpu。训练开始后你会在终端看到损失值逐渐下降。DETR的损失包括分类损失、L1边界框损失和GIoU损失。训练初期损失可能波动较大这是正常的。程序默认每10个epoch会在验证集上评估一次输出AP平均精度指标。这里有个非常重要的点DETR训练收敛相对较慢你可能在前50个epoch都看不到明显的AP提升甚至mAP还是0请不要灰心这是DETR训练的特性只要损失在稳步下降就让它继续跑。如果你的训练中断了想从上次保存的检查点继续训练可以这样python main.py \ --dataset_file coco \ --coco_path /path/to/your/coco_custom \ --output_dir /path/to/save/checkpoints \ --resume /path/to/save/checkpoints/checkpoint.pth \ # 指定具体的检查点文件 --start_epoch 100 \ # 指定从第几轮开始查看日志确认 --epochs 3007. 模型验证与可视化看看你的DETR学会了什么训练完成后或者训练中途你想看看模型效果就需要进行推理验证。官方代码的评估是在验证集上计算mAP但我们通常更想直观地看几张图片的检测效果。我写了一个简单的推理脚本inference_demo.py它可以加载训练好的模型对单张或一批图片进行预测并把结果画出来。这个脚本的核心步骤是1加载模型和权重2对输入图片进行预处理缩放、归一化3模型前向传播4根据置信度阈值过滤预测结果5将归一化的边界框坐标转换回原图尺寸6用OpenCV或Matplotlib画框并保存。import torch import torchvision.transforms as T import numpy as np from PIL import Image import cv2 import os import argparse from models.detr import build_detr_model # 假设你有一个构建模型的函数 # 1. 定义预处理变换必须与训练时一致 transform T.Compose([ T.Resize(800), # 将短边缩放到800长边按比例缩放 T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # ImageNet均值标准差 ]) # 2. 定义你的类别列表必须与训练时顺序一致 CLASSES [ i2, i4, i5, il100, il60, il80, io, ip, p10, p11, p12, p19, p23, p26, p27, p3, p5, p6, pg, ph4, ph4d5, ph5, pl100, pl120, pl20, pl30, pl40, pl5, pl50, pl60, pl70, pl80, pm20, pm30, pm55, pn, pne, po, pr40, w13, w32, w55, w57, w59, wo, crosswalk, left, right, straight, straightl, straightr, diamond, noparking, person, car, truck, bus ] # 3. 坐标转换函数模型输出是cxcywh归一化格式 def box_cxcywh_to_xyxy(x): 将 [center_x, center_y, width, height] 转换为 [x_min, y_min, x_max, y_max] x_c, y_c, w, h x.unbind(1) b [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c 0.5 * w), (y_c 0.5 * h)] return torch.stack(b, dim1) def rescale_bboxes(out_bbox, img_size): 将归一化的边界框坐标缩放到原图尺寸 img_w, img_h img_size b box_cxcywh_to_xyxy(out_bbox) b b * torch.tensor([img_w, img_h, img_w, img_h], dtypetorch.float32) return b # 4. 绘图函数 def plot_results(pil_img, scores, boxes, labels, save_path, threshold0.7): 将预测结果绘制到图片上并保存。 pil_img: PIL Image对象 scores: 置信度张量 boxes: 边界框张量格式 [x_min, y_min, x_max, y_max]绝对坐标 labels: 类别ID张量 save_path: 保存路径 threshold: 置信度阈值 img np.array(pil_img) # 为不同类别生成一些颜色 colors plt.cm.hsv(np.linspace(0, 1, len(CLASSES)1)).tolist() keep scores threshold scores scores[keep] boxes boxes[keep] labels labels[keep] for score, box, label_id in zip(scores, boxes, labels): box [int(coord) for coord in box] label CLASSES[label_id] color colors[label_id] color [int(c*255) for c in color[:3]] # 转换为0-255的BGR color tuple(color[::-1]) # RGB - BGR for OpenCV # 画矩形框 cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), color, 2) # 准备标签文本 text f{label}: {score:.2f} # 获取文本大小 (text_width, text_height), baseline cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) # 画文本背景框 cv2.rectangle(img, (box[0], box[1] - text_height - baseline), (box[0] text_width, box[1]), color, -1) # 写文本 cv2.putText(img, text, (box[0], box[1] - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) # 保存图片 cv2.imwrite(save_path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR)) print(f结果已保存至: {save_path}) def main(args): # 5. 加载模型 print(正在加载模型...) model, criterion, postprocessors build_detr_model(args) # 你需要根据你的模型构建方式调整 checkpoint torch.load(args.checkpoint, map_locationcpu) model.load_state_dict(checkpoint[model]) model.to(args.device) model.eval() print(模型加载完毕。) # 6. 遍历图片进行推理 image_extensions (.jpg, .jpeg, .png, .bmp) image_files [f for f in os.listdir(args.image_dir) if f.lower().endswith(image_extensions)] for img_file in image_files: img_path os.path.join(args.image_dir, img_file) # 读取图片 pil_img Image.open(img_path).convert(RGB) original_size pil_img.size # (width, height) # 预处理 img_tensor transform(pil_img).unsqueeze(0).to(args.device) # 增加batch维度 # 推理 with torch.no_grad(): outputs model(img_tensor) # 后处理这里简化处理直接取最后一层的输出 # DETR输出格式outputs[pred_logits], outputs[pred_boxes] probas outputs[pred_logits].softmax(-1)[0, :, :-1] # 去掉“无物体”类 keep probas.max(-1).values args.conf_thresh # 获取类别和分数 scores, labels probas[keep].max(-1) # 获取边界框并缩放 bboxes rescale_bboxes(outputs[pred_boxes][0, keep], original_size) # 保存结果 save_name os.path.splitext(img_file)[0] _det.jpg save_path os.path.join(args.output_dir, save_name) plot_results(pil_img, scores, bboxes, labels, save_path, thresholdargs.conf_thresh) if __name__ __main__: parser argparse.ArgumentParser(DETR推理演示) parser.add_argument(--checkpoint, requiredTrue, help训练好的模型权重路径 (.pth)) parser.add_argument(--image-dir, requiredTrue, help待检测图片文件夹路径) parser.add_argument(--output-dir, default./infer_results, help检测结果保存路径) parser.add_argument(--conf-thresh, typefloat, default0.7, help置信度阈值) parser.add_argument(--device, defaultcuda if torch.cuda.is_available() else cpu, help计算设备) # 可能需要添加模型构建所需的参数这里省略你需要根据你的模型定义补充 args parser.parse_args() # 创建输出目录 os.makedirs(args.output_dir, exist_okTrue) main(args)运行这个脚本你就能在output_dir里看到画好预测框的图片了。通过调整--conf-thresh参数可以控制只显示高置信度的预测结果这对于分析模型性能非常有用。第一次看到自己训练的模型在全新的图片上正确框出目标时那种成就感是非常棒的。这个过程里你可能需要反复调整数据质量、模型参数甚至回去检查标注但每一次迭代都会让你对DETR和整个目标检测流程有更深的理解。