1.DETR介绍

DETR(Detection with TRansformers)是基于transformer的端对端目标检测，无NMS后处理步骤，无anchor。
代码链接：https://github.com/facebookresearch/detr

2.数据集处理

DETR需要的数据集格式为coco格式，这里我是用自己的YOLO格式数据集转化成COCO格式，然后进行训练的。
YOLO数据集的组织格式是：
其中images里面分别存放训练集train和验证集val的图片，labels存放训练集train和验证集val的txt标签。

要转化成适应DETR模型读取的COCO数据集的组织形式是：
其中train2017存放训练集的图片，val2017存放验证集的图片，
annotations文件夹里面存放train和val的json标签。

下面是转化代码：

• 需要进行类别映射，每个类别对应的id分别存放在categories里面，这里我没有用classes.txt文件存放，相当于直接把classes.txt里面的类别写出来了。
• 我的图片是png格式的，如果图片是jpg格式的，将png改成jpg即可。image_name = filename.replace(‘.txt’, ‘.jpg’)
• 最后修改文件路径，改成自己的路径，这里最后会输出train和val的json文件，图片不会处理，按上述目录组织形式将图片组织起来即可。
• 生成的文件夹记得改为instances_train2017.json这种样子

import os
import json
from PIL import Image# 定义类别映射
categories = [{"id": 0, "name": "Double hexagonal column"},{"id": 1, "name": "Flange nut"},{"id": 2, "name": "Hexagon nut"},{"id": 3, "name": "Hexagon pillar"},{"id": 4, "name": "Hexagon screw"},{"id": 5, "name": "Hexagonal steel column"},{"id": 6, "name": "Horizontal bubble"},{"id": 7, "name": "Keybar"},{"id": 8, "name": "Plastic cushion pillar"},{"id": 9, "name": "Rectangular nut"},{"id": 10, "name": "Round head screw"},{"id": 11, "name": "Spring washer"},{"id": 12, "name": "T-shaped screw"}
]def yolo_to_coco(yolo_images_dir, yolo_labels_dir, output_json_path):# 初始化 COCO 数据结构data = {"images": [],"annotations": [],"categories": categories}image_id = 1annotation_id = 1def get_image_size(image_path):with Image.open(image_path) as img:return img.width, img.height# 遍历标签目录for filename in os.listdir(yolo_labels_dir):if not filename.endswith('.txt'):continue  # 只处理 .txt 文件image_name = filename.replace('.txt', '.png')# 如果图片是jpg格式的，将png改成jpg即可。image_path = os.path.join(yolo_images_dir, image_name)if not os.path.exists(image_path):print(f"⚠️ 警告: 图像 {image_name} 不存在，跳过 {filename}")continueimage_width, image_height = get_image_size(image_path)image_info = {"id": image_id,"width": image_width,"height": image_height,"file_name": image_name}data["images"].append(image_info)with open(os.path.join(yolo_labels_dir, filename), 'r') as file:lines = file.readlines()for line in lines:parts = line.strip().split()if len(parts) != 5:print(f"⚠️ 警告: 标签 {filename} 格式错误: {line.strip()}")continuecategory_id = int(parts[0])x_center = float(parts[1]) * image_widthy_center = float(parts[2]) * image_heightbbox_width = float(parts[3]) * image_widthbbox_height = float(parts[4]) * image_heightx_min = int(x_center - bbox_width / 2)y_min = int(y_center - bbox_height / 2)bbox = [x_min, y_min, bbox_width, bbox_height]area = bbox_width * bbox_heightannotation_info = {"id": annotation_id,"image_id": image_id,"category_id": category_id,"bbox": bbox,"area": area,"iscrowd": 0}data["annotations"].append(annotation_info)annotation_id += 1image_id += 1os.makedirs(os.path.dirname(output_json_path), exist_ok=True)with open(output_json_path, 'w') as json_file:json.dump(data, json_file, indent=4)print(f"✅ 转换完成: {output_json_path}")# 输入路径 (YOLO 格式数据集)
yolo_base_dir = "/home/yu/Yolov8/ultralytics-main/mydata0"
yolo_train_images = os.path.join(yolo_base_dir, "images/train")
yolo_train_labels = os.path.join(yolo_base_dir, "labels/train")
yolo_val_images = os.path.join(yolo_base_dir, "images/val")
yolo_val_labels = os.path.join(yolo_base_dir, "labels/val")# 输出路径 (COCO 格式)
coco_base_dir = "/home/yu/Yolov8/ultralytics-main/mydata0_coco"
coco_train_json = os.path.join(coco_base_dir, "annotations/instances_train.json")
coco_val_json = os.path.join(coco_base_dir, "annotations/instances_val.json")# 运行转换
yolo_to_coco(yolo_train_images, yolo_train_labels, coco_train_json)
yolo_to_coco(yolo_val_images, yolo_val_labels, coco_val_json)

3.转化结果可视化

COCO数据集JSON文件格式分为以下几个字段。

{"info": info, # dict"licenses": [license], # list ，内部是dict"images": [image], # list ，内部是dict"annotations": [annotation], # list ，内部是dict"categories": # list ，内部是dict}

可以运行以下脚本查看转化后的标签是否与图片目标对应：

• 修改代码的json_path和img_path，json_path是标签对应的路径，img_path是图像对应的路径

'''
该代码的功能是：读取图像以及对应bbox的信息
'''
import os
from pycocotools.coco import COCO
from PIL import Image, ImageDraw
import matplotlib.pyplot as pltjson_path = "/home/yu/Yolov8/ultralytics-main/mydata0_coco/annotations/instances_val.json"
img_path = ("/home/yu/Yolov8/ultralytics-main/mydata0_coco/images/val")# load coco data
coco = COCO(annotation_file=json_path)# get all image index info
ids = list(sorted(coco.imgs.keys()))
print("number of images: {}".format(len(ids)))# get all coco class labels
coco_classes = dict([(v["id"], v["name"]) for k, v in coco.cats.items()])# 遍历前三张图像
for img_id in ids[:3]:# 获取对应图像id的所有annotations idx信息ann_ids = coco.getAnnIds(imgIds=img_id)# 根据annotations idx信息获取所有标注信息targets = coco.loadAnns(ann_ids)# get image file namepath = coco.loadImgs(img_id)[0]['file_name']# read imageimg = Image.open(os.path.join(img_path, path)).convert('RGB')draw = ImageDraw.Draw(img)# draw box to imagefor target in targets:x, y, w, h = target["bbox"]x1, y1, x2, y2 = x, y, int(x + w), int(y + h)draw.rectangle((x1, y1, x2, y2))draw.text((x1, y1), coco_classes[target["category_id"]])# show imageplt.imshow(img)plt.show()

运行该代码，你将会看到你的标签是否对应：
如果目标没有边界框则说明你转化的json不对！

4.数据集训练

4.1修改pth文件

将它的pth文件改一下，因为他是用的coco数据集，而我们只需要训练自己的数据集，就是下图这个文件，这是它原本的

新建一个.py文件，运行下面代码，就会生成一个你数据集所需要的物体数目的pth，记得改类别数！。

import torch
pretrained_weights  = torch.load('detr-r50-e632da11.pth')num_class = 14 #这里是你的物体数+1，因为背景也算一个
pretrained_weights["model"]["class_embed.weight"].resize_(num_class+1, 256)
pretrained_weights["model"]["class_embed.bias"].resize_(num_class+1)
torch.save(pretrained_weights, "detr-r50_%d.pth"%num_class

这是我们生成的。

4.2类别参数修改

修改models/detr.py文件，build()函数中，可以将红框部分的代码都注释掉，直接设置num_classes为自己的类别数+1
因为我的类别数是13，所以我这里num_classes=14

4.3训练

修改main.py文件的epochs、lr、batch_size等训练参数：
以下这些参数都在get_args_parser()函数里面。

修改自己的数据集路径：

设置输出路径：

修改resume为自己的预训练权重文件路径
这里就是你刚才运行脚本生成的pth文件的路径：

运行main.py文件
或者可以通过命令行运行：

python main.py --dataset_file "coco" --coco_path "/home/yu/Yolov8/ultralytics-main/mydata0_coco" --epoch 300 --lr=1e-4 --batch_size=8 --num_workers=4 --output_dir="outputs" --resume="detr_r50_14.pth"

5.成功运行！

6.参考文献

1.【DETR】训练自己的数据集-实践笔记
2. yolo数据集格式（txt）转coco格式，方便mmyolo转标签格式
3. windows10复现DEtection TRansformers（DETR）并实现自己的数据集