【DETR实战】从零构建遥感影像检测模型：NWPU VHR-10数据集全流程解析

1. 从零开始DETR与遥感影像检测的那些事儿第一次接触DETR框架时我被它简洁优雅的设计惊艳到了。这个基于Transformer的目标检测模型完全摒弃了传统方法中复杂的anchor设置和后处理NMS步骤让整个检测流程变得前所未有的清晰。特别是在处理遥感影像这种特殊场景时DETR展现出了独特的优势。NWPU VHR-10数据集是遥感领域非常经典的基准数据集包含10类常见地物目标从飞机、船舶到运动场地应有尽有。但实际操作时会发现这个数据集与DETR要求的COCO格式存在不小差距。我刚开始就踩了个坑——直接拿着原始数据就开始训练结果模型完全无法收敛。后来才发现格式转换这一步绝对不能马虎。遥感影像检测最大的特点就是目标尺度变化大。同一张图片里可能既有占据几十个像素的小型车辆又有横跨整幅图像的大型机场跑道。DETR的全局注意力机制在这里派上了大用场它能同时捕捉不同尺度的目标特征这是传统CNN-based检测器难以做到的。2. 数据准备从原始数据到DETR可用的COCO格式2.1 理解NWPU数据集的原始结构NWPU VHR-10的原始数据组织方式很有特点。它分为正样本集650张和负样本集150张我们主要使用正样本集。每张图片对应一个txt格式的标注文件标注格式是这样的(xmin,ymin)(xmax,ymax) category_id。这种格式和常见的VOC或COCO都不太一样需要专门处理。我建议先花点时间浏览数据集。用Python的PIL库简单写个可视化脚本把标注框画在图片上检查一下from PIL import Image, ImageDraw import os image_path NWPU/positive image set/100.jpg label_path NWPU/ground truth/100.txt img Image.open(image_path) draw ImageDraw.Draw(img) with open(label_path) as f: for line in f: if line.strip(): coords line.strip().split(,) xmin int(coords[0].split(()[1]) ymin int(coords[1].split())[0]) xmax int(coords[2].split(()[1]) ymax int(coords[3].split())[0]) draw.rectangle([xmin, ymin, xmax, ymax], outlinered, width2) img.show()2.2 构建COCO格式的标注文件DETR要求的数据格式是COCO风格的这意味着我们需要把原始数据转换成包含images、annotations、categories三个关键字段的JSON文件。我写了一个转换脚本核心逻辑是这样的遍历所有图片为每张图片生成唯一的image_id解析对应的标注文件为每个目标实例生成包含bbox、category_id等信息的annotation按照COCO的格式要求组装成完整的字典结构关键点在于area字段的计算这个在原始数据中是没有的。area就是bbox的宽度乘以高度但很多新手会忽略这个细节导致训练时报KeyError。另外category_id需要从0开始还是从1开始也要特别注意这关系到后续模型输出的类别对应关系。def convert_nwpu_to_coco(image_dir, label_dir, output_json): images [] annotations [] categories [{id: i1, name: name} for i, name in enumerate(NWPU_CATEGORIES)] image_id 1 annotation_id 1 for img_name in os.listdir(image_dir): # 处理图片信息 img_path os.path.join(image_dir, img_name) width, height Image.open(img_path).size images.append({ id: image_id, file_name: img_name, width: width, height: height }) # 处理标注信息 label_path os.path.join(label_dir, img_name.replace(.jpg, .txt)) with open(label_path) as f: for line in f: if line.strip(): parts line.strip().split(,) xmin float(parts[0].split(()[1]) ymin float(parts[1].split())[0]) xmax float(parts[2].split(()[1]) ymax float(parts[3].split())[0]) category_id int(parts[4]) width xmax - xmin height ymax - ymin area width * height annotations.append({ id: annotation_id, image_id: image_id, category_id: category_id, bbox: [xmin, ymin, width, height], area: area, iscrowd: 0 }) annotation_id 1 image_id 1 coco_dict { images: images, annotations: annotations, categories: categories } with open(output_json, w) as f: json.dump(coco_dict, f)3. 模型配置让DETR适配遥感场景3.1 修改预训练权重DETR官方提供的预训练模型是在COCO数据集上训练的有91个类别包括背景。我们的NWPU数据集只有10个类别需要调整分类头的维度。这个步骤很关键直接影响到模型能否正常训练。import torch # 加载官方预训练权重 pretrained torch.load(detr-r50-e632da11.pth) # 修改分类头参数 num_classes 10 # NWPU有10类 pretrained[model][class_embed.weight].resize_(num_classes1, 256) pretrained[model][class_embed.bias].resize_(num_classes1) # 保存修改后的权重 torch.save(pretrained, detr-r50-nwpu.pth)3.2 调整模型参数遥感影像中的目标通常比较小我建议做以下调整增加decoder的层数增强模型对小目标的检测能力调整学习率遥感数据通常需要更小的学习率增加训练epoch因为遥感数据的分布和自然图像有差异在main.py中可以这样修改默认参数def get_args_parser(): parser argparse.ArgumentParser(Set transformer detector, add_helpFalse) parser.add_argument(--lr, default1e-5, typefloat) # 降低学习率 parser.add_argument(--lr_backbone, default1e-6, typefloat) parser.add_argument(--epochs, default500, typeint) # 增加epoch parser.add_argument(--dec_layers, default6, typeint) # 增加decoder层数 return parser4. 训练技巧提升遥感检测效果的关键4.1 数据增强策略针对遥感影像的特点我推荐使用以下增强组合随机裁剪模拟不同尺度的目标颜色抖动补偿不同成像条件带来的颜色差异随机旋转遥感影像通常没有固定的方向性可以在datasets/coco.py中修改增强管道def make_coco_transforms(image_set): normalize T.Compose([ T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) if image_set train: return T.Compose([ T.RandomHorizontalFlip(), T.RandomVerticalFlip(), # 新增垂直翻转 T.RandomRotation(30), # 随机旋转 T.ColorJitter(brightness0.3, contrast0.3, saturation0.3), # 颜色抖动 T.RandomResize([480, 512, 544, 576, 608], max_size800), normalize, ]) return T.Compose([ T.RandomResize([800], max_size1333), normalize, ])4.2 损失函数调优DETR默认使用匈牙利匹配算法计算损失但对于遥感场景可能需要调整以下几点分类损失和bbox损失的权重比例GIoU损失的权重辅助解码层的损失权重在models/detr.py中可以这样调整class DETR(nn.Module): def __init__(self, ...): # 修改损失权重 matcher HungarianMatcher( cost_class2, # 提高分类损失的权重 cost_bbox5, cost_giou2 ) losses [labels, boxes, cardinality] weight_dict { loss_ce: 2, # 分类损失 loss_bbox: 5, # bbox回归损失 loss_giou: 2 # GIoU损失 }5. 评估与可视化解读模型表现5.1 评估指标分析训练完成后DETR会自动输出COCO风格的评估指标包括AP、AP50、AP75等。对于遥感影像我特别关注小目标的APAPs中等目标的APAPm不同类别的单独AP值可以通过修改main.py中的评估代码来获取更详细的信息from pycocotools.coco import COCO from pycocotools.cocoeval import COCOeval def evaluate(...): # 在原有评估基础上增加小目标分析 coco_eval COCOeval(coco_gt, coco_dt, iouType) coco_eval.evaluate() coco_eval.accumulate() coco_eval.summarize() # 打印小目标结果 stats coco_eval.stats print(fSmall objects AP: {stats[3]:.3f}) print(fMedium objects AP: {stats[4]:.3f}) print(fLarge objects AP: {stats[5]:.3f})5.2 结果可视化技巧DETR的预测结果可视化很有特点因为它是并行输出所有预测框。我推荐使用以下方式增强可视化效果为不同类别使用不同颜色显示预测分数添加注意力热图这里有个实用的可视化函数def visualize_prediction(image_path, model, transform, device): img Image.open(image_path) img_tensor transform(img).unsqueeze(0).to(device) with torch.no_grad(): outputs model(img_tensor) # 转换到原图尺寸 orig_size torch.as_tensor([img.size[::-1]]).to(device) results postprocess(outputs, orig_size)[0] draw ImageDraw.Draw(img) colors [red, blue, green, yellow, purple, orange, pink, cyan, magenta, lime] for score, label, box in zip(results[scores], results[labels], results[boxes]): if score 0.7: # 只显示高置信度的预测 box box.cpu().numpy() draw.rectangle( [(box[0], box[1]), (box[2], box[3])], outlinecolors[label % len(colors)], width3 ) draw.text( (box[0], box[1]), f{NWPU_CATEGORIES[label]}: {score:.2f}, fillcolors[label % len(colors)] ) return img在完成整个流程后我发现DETR在遥感影像检测上的表现确实令人惊喜。虽然训练初期收敛较慢但一旦过了前50个epoch性能提升非常明显。特别是对于密集小目标的检测DETR的全局注意力机制展现出了比传统方法更好的性能。不过要注意的是遥感影像通常分辨率很高训练时需要适当调整输入尺寸否则显存很容易爆掉。