【DETR】训练自己的数据集以及YOLO数据集格式(txt)转化成COCO格式（json） ...

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这种样子
  

1. import os
2. import json
3. from PIL import Image
5. # 定义类别映射
6. categories = [
7.     {"id": 0, "name": "Double hexagonal column"},
8.     {"id": 1, "name": "Flange nut"},
9.     {"id": 2, "name": "Hexagon nut"},
10.     {"id": 3, "name": "Hexagon pillar"},
11.     {"id": 4, "name": "Hexagon screw"},
12.     {"id": 5, "name": "Hexagonal steel column"},
13.     {"id": 6, "name": "Horizontal bubble"},
14.     {"id": 7, "name": "Keybar"},
15.     {"id": 8, "name": "Plastic cushion pillar"},
16.     {"id": 9, "name": "Rectangular nut"},
17.     {"id": 10, "name": "Round head screw"},
18.     {"id": 11, "name": "Spring washer"},
19.     {"id": 12, "name": "T-shaped screw"}
20. ]
22. def yolo_to_coco(yolo_images_dir, yolo_labels_dir, output_json_path):
23.     # 初始化 COCO 数据结构
24.     data = {
25.         "images": [],
26.         "annotations": [],
27.         "categories": categories
28.     }
30.     image_id = 1
31.     annotation_id = 1
33.     def get_image_size(image_path):
34.         with Image.open(image_path) as img:
35.             return img.width, img.height
37.     # 遍历标签目录
38.     for filename in os.listdir(yolo_labels_dir):
39.         if not filename.endswith('.txt'):
40.             continue  # 只处理 .txt 文件
42.         image_name = filename.replace('.txt', '.png')# 如果图片是jpg格式的，将png改成jpg即可。
43.         
44.         image_path = os.path.join(yolo_images_dir, image_name)
46.         if not os.path.exists(image_path):
47.             print(f"⚠️ 警告: 图像 {image_name} 不存在，跳过 {filename}")
48.             continue
50.         image_width, image_height = get_image_size(image_path)
52.         image_info = {
53.             "id": image_id,
54.             "width": image_width,
55.             "height": image_height,
56.             "file_name": image_name
57.         }
58.         data["images"].append(image_info)
60.         with open(os.path.join(yolo_labels_dir, filename), 'r') as file:
61.             lines = file.readlines()
63.         for line in lines:
64.             parts = line.strip().split()
65.             if len(parts) != 5:
66.                 print(f"⚠️ 警告: 标签 {filename} 格式错误: {line.strip()}")
67.                 continue
69.             category_id = int(parts[0])
70.             x_center = float(parts[1]) * image_width
71.             y_center = float(parts[2]) * image_height
72.             bbox_width = float(parts[3]) * image_width
73.             bbox_height = float(parts[4]) * image_height
75.             x_min = int(x_center - bbox_width / 2)
76.             y_min = int(y_center - bbox_height / 2)
77.             bbox = [x_min, y_min, bbox_width, bbox_height]
78.             area = bbox_width * bbox_height
80.             annotation_info = {
81.                 "id": annotation_id,
82.                 "image_id": image_id,
83.                 "category_id": category_id,
84.                 "bbox": bbox,
85.                 "area": area,
86.                 "iscrowd": 0
87.             }
88.             data["annotations"].append(annotation_info)
89.             annotation_id += 1
91.         image_id += 1
93.     os.makedirs(os.path.dirname(output_json_path), exist_ok=True)
94.     with open(output_json_path, 'w') as json_file:
95.         json.dump(data, json_file, indent=4)
97.     print(f"✅ 转换完成: {output_json_path}")
100. # 输入路径 (YOLO 格式数据集)
101. yolo_base_dir = "/home/yu/Yolov8/ultralytics-main/mydata0"
102. yolo_train_images = os.path.join(yolo_base_dir, "images/train")
103. yolo_train_labels = os.path.join(yolo_base_dir, "labels/train")
104. yolo_val_images = os.path.join(yolo_base_dir, "images/val")
105. yolo_val_labels = os.path.join(yolo_base_dir, "labels/val")
107. # 输出路径 (COCO 格式)
108. coco_base_dir = "/home/yu/Yolov8/ultralytics-main/mydata0_coco"
109. coco_train_json = os.path.join(coco_base_dir, "annotations/instances_train.json")
110. coco_val_json = os.path.join(coco_base_dir, "annotations/instances_val.json")
112. # 运行转换
113. yolo_to_coco(yolo_train_images, yolo_train_labels, coco_train_json)
114. yolo_to_coco(yolo_val_images, yolo_val_labels, coco_val_json)

复制代码

3.转化效果可视化

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

1. {
2.     "info": info, # dict
3.      "licenses": [license], # list ，内部是dict
4.      "images": [image], # list ，内部是dict
5.      "annotations": [annotation], # list ，内部是dict
6.      "categories": # list ，内部是dict
7. }

复制代码

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

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

1. '''
2. 该代码的功能是：读取图像以及对应bbox的信息
3. '''
4. import os
5. from pycocotools.coco import COCO
6. from PIL import Image, ImageDraw
7. import matplotlib.pyplot as plt
9. json_path = "/home/yu/Yolov8/ultralytics-main/mydata0_coco/annotations/instances_val.json"
10. img_path = ("/home/yu/Yolov8/ultralytics-main/mydata0_coco/images/val")
12. # load coco data
13. coco = COCO(annotation_file=json_path)
15. # get all image index info
16. ids = list(sorted(coco.imgs.keys()))
17. print("number of images: {}".format(len(ids)))
19. # get all coco class labels
20. coco_classes = dict([(v["id"], v["name"]) for k, v in coco.cats.items()])
22. # 遍历前三张图像
23. for img_id in ids[:3]:
24.     # 获取对应图像id的所有annotations idx信息
25.     ann_ids = coco.getAnnIds(imgIds=img_id)
27.     # 根据annotations idx信息获取所有标注信息
28.     targets = coco.loadAnns(ann_ids)
30.     # get image file name
31.     path = coco.loadImgs(img_id)[0]['file_name']
33.     # read image
34.     img = Image.open(os.path.join(img_path, path)).convert('RGB')
35.     draw = ImageDraw.Draw(img)
36.     # draw box to image
37.     for target in targets:
38.         x, y, w, h = target["bbox"]
39.         x1, y1, x2, y2 = x, y, int(x + w), int(y + h)
40.         draw.rectangle((x1, y1, x2, y2))
41.         draw.text((x1, y1), coco_classes[target["category_id"]])
43.     # show image
44.     plt.imshow(img)
45.     plt.show()

复制代码

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

4.数据集训练

4.1修改pth文件

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

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

1. import torch
2. pretrained_weights  = torch.load('detr-r50-e632da11.pth')
4. num_class = 14 #这里是你的物体数+1，因为背景也算一个
5. pretrained_weights["model"]["class_embed.weight"].resize_(num_class+1, 256)
6. pretrained_weights["model"]["class_embed.bias"].resize_(num_class+1)
7. 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文件
大概可以通过命令行运行：

1. 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）并实现自己的数据集

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。