python：VOC格式数据集转换为YOLO数据集格式_voc格式转化为yolo

作者：CSDN @ _养乐多_

本文将介绍如何将目标检测中常用的VOC格式数据集转换为YOLO数据集，并进行数据集比例划分，从而方便的进行YOLO目标检测。

如果不想分两步，可以直接看第三节代码。

一、将VOC格式数据集转换为YOLO格式数据集

执行以下脚本将VOC格式数据集转换为YOLO格式数据集。
但是需要注意的是：

1. 转换之后的数据集只有Images和labels两个文件。还需要执行第二节中的脚本进行数据集划分，将总的数据集划分为训练、验证、测试数据集；
2. 使用的话，需要修改 class_mapping 中类别名和对应标签，还有VOC数据集路径、YOLO数据集路径。

import os
import shutil
import xml.etree.ElementTree as ET

# VOC格式数据集路径
voc_data_path = 'E:\\DataSet\\helmet-VOC'
voc_annotations_path = os.path.join(voc_data_path, 'Annotations')
voc_images_path = os.path.join(voc_data_path, 'JPEGImages')

# YOLO格式数据集保存路径
yolo_data_path = 'E:\\DataSet\\helmet-YOLO'
yolo_images_path = os.path.join(yolo_data_path, 'images')
yolo_labels_path = os.path.join(yolo_data_path, 'labels')

# 创建YOLO格式数据集目录
os.makedirs(yolo_images_path, exist_ok=True)
os.makedirs(yolo_labels_path, exist_ok=True)

# 类别映射 (可以根据自己的数据集进行调整)
class_mapping = {
    'head': 0,
    'helmet': 1,
    'person': 2,
    # 添加更多类别...
}

def convert_voc_to_yolo(voc_annotation_file, yolo_label_file):
    tree = ET.parse(voc_annotation_file)
    root = tree.getroot()

    size = root.find('size')
    width = float(size.find('width').text)
    height = float(size.find('height').text)

    with open(yolo_label_file, 'w') as f:
        for obj in root.findall('object'):
            cls = obj.find('name').text
            if cls not in class_mapping:
                continue
            cls_id = class_mapping[cls]
            xmlbox = obj.find('bndbox')
            xmin = float(xmlbox.find('xmin').text)
            ymin = float(xmlbox.find('ymin').text)
            xmax = float(xmlbox.find('xmax').text)
            ymax = float(xmlbox.find('ymax').text)

            x_center = (xmin + xmax) / 2.0 / width
            y_center = (ymin + ymax) / 2.0 / height
            w = (xmax - xmin) / width
            h = (ymax - ymin) / height

            f.write(f"{cls_id} {x_center} {y_center} {w} {h}\n")

# 遍历VOC数据集的Annotations目录，进行转换
for voc_annotation in os.listdir(voc_annotations_path):
    if voc_annotation.endswith('.xml'):
        voc_annotation_file = os.path.join(voc_annotations_path, voc_annotation)
        image_id = os.path.splitext(voc_annotation)[0]
        voc_image_file = os.path.join(voc_images_path, f"{image_id}.jpg")
        yolo_label_file = os.path.join(yolo_labels_path, f"{image_id}.txt")
        yolo_image_file = os.path.join(yolo_images_path, f"{image_id}.jpg")

        convert_voc_to_yolo(voc_annotation_file, yolo_label_file)
        if os.path.exists(voc_image_file):
            shutil.copy(voc_image_file, yolo_image_file)

print("转换完成！")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

二、YOLO格式数据集划分（训练、验证、测试）

随机将数据集按照0.7-0.2-0.1比例划分为训练、验证、测试数据集。
注意，修改代码中图片的后缀，如果是.jpg，就把.png修改为.jpg。

最终结果，

import os
import shutil
import random

# YOLO格式数据集保存路径
yolo_images_path1 = 'E:\\DataSet\\helmet-VOC'
yolo_labels_path1 = 'E:\\DataSet\\helmet-YOLO'
yolo_data_path = yolo_labels_path1

yolo_images_path = os.path.join(yolo_images_path1, 'JPEGImages')
yolo_labels_path = os.path.join(yolo_labels_path1, 'labels')

# 创建划分后的目录结构
train_images_path = os.path.join(yolo_data_path, 'train', 'images')
train_labels_path = os.path.join(yolo_data_path, 'train', 'labels')
val_images_path = os.path.join(yolo_data_path, 'val', 'images')
val_labels_path = os.path.join(yolo_data_path, 'val', 'labels')
test_images_path = os.path.join(yolo_data_path, 'test', 'images')
test_labels_path = os.path.join(yolo_data_path, 'test', 'labels')

os.makedirs(train_images_path, exist_ok=True)
os.makedirs(train_labels_path, exist_ok=True)
os.makedirs(val_images_path, exist_ok=True)
os.makedirs(val_labels_path, exist_ok=True)
os.makedirs(test_images_path, exist_ok=True)
os.makedirs(test_labels_path, exist_ok=True)

# 获取所有图片文件名（不包含扩展名）
image_files = [f[:-4] for f in os.listdir(yolo_images_path) if f.endswith('.png')]

# 随机打乱文件顺序
random.shuffle(image_files)

# 划分数据集比例
train_ratio = 0.7
val_ratio = 0.2
test_ratio = 0.1

train_count = int(train_ratio * len(image_files))
val_count = int(val_ratio * len(image_files))
test_count = len(image_files) - train_count - val_count

train_files = image_files[:train_count]
val_files = image_files[train_count:train_count + val_count]
test_files = image_files[train_count + val_count:]

# 移动文件到相应的目录
def move_files(files, src_images_path, src_labels_path, dst_images_path, dst_labels_path):
    for file in files:
        src_image_file = os.path.join(src_images_path, f"{file}.png")
        src_label_file = os.path.join(src_labels_path, f"{file}.txt")
        dst_image_file = os.path.join(dst_images_path, f"{file}.png")
        dst_label_file = os.path.join(dst_labels_path, f"{file}.txt")

        if os.path.exists(src_image_file) and os.path.exists(src_label_file):
            shutil.move(src_image_file, dst_image_file)
            shutil.move(src_label_file, dst_label_file)

# 移动训练集文件
move_files(train_files, yolo_images_path, yolo_labels_path, train_images_path, train_labels_path)
# 移动验证集文件
move_files(val_files, yolo_images_path, yolo_labels_path, val_images_path, val_labels_path)
# 移动测试集文件
move_files(test_files, yolo_images_path, yolo_labels_path, test_images_path, test_labels_path)

print("数据集划分完成！")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

三、一步到位

如果不想分两步进行格式转换，那么以下脚本结合了以上两步，直接得到最后按比例划分训练、验证、测试的数据集结果。

注意：需要修改 voc_data_path ，yolo_data_path ，class_mapping 以及 ‘.png’ 后缀。

import os
import shutil
import random
import xml.etree.ElementTree as ET
from tqdm import tqdm

# VOC格式数据集路径
voc_data_path = 'E:\\DataSet-VOC'
voc_annotations_path = os.path.join(voc_data_path, 'Annotations')
voc_images_path = os.path.join(voc_data_path, 'JPEGImages')

# YOLO格式数据集保存路径
yolo_data_path = 'E:\\DataSet-YOLO'
yolo_images_path = os.path.join(yolo_data_path, 'images')
yolo_labels_path = os.path.join(yolo_data_path, 'labels')

# 创建YOLO格式数据集目录
os.makedirs(yolo_images_path, exist_ok=True)
os.makedirs(yolo_labels_path, exist_ok=True)

# 类别映射 (可以根据自己的数据集进行调整)
class_mapping = {
    'head': 0,
    'helmet': 1,
    'person': 2,
    # 添加更多类别...
}

def convert_voc_to_yolo(voc_annotation_file, yolo_label_file):
    tree = ET.parse(voc_annotation_file)
    root = tree.getroot()

    size = root.find('size')
    width = float(size.find('width').text)
    height = float(size.find('height').text)

    with open(yolo_label_file, 'w') as f:
        for obj in root.findall('object'):
            cls = obj.find('name').text
            if cls not in class_mapping:
                continue
            cls_id = class_mapping[cls]
            xmlbox = obj.find('bndbox')
            xmin = float(xmlbox.find('xmin').text)
            ymin = float(xmlbox.find('ymin').text)
            xmax = float(xmlbox.find('xmax').text)
            ymax = float(xmlbox.find('ymax').text)

            x_center = (xmin + xmax) / 2.0 / width
            y_center = (ymin + ymax) / 2.0 / height
            w = (xmax - xmin) / width
            h = (ymax - ymin) / height

            f.write(f"{cls_id} {x_center} {y_center} {w} {h}\n")

# 遍历VOC数据集的Annotations目录，进行转换
print("开始VOC到YOLO格式转换...")
for voc_annotation in tqdm(os.listdir(voc_annotations_path)):
    if voc_annotation.endswith('.xml'):
        voc_annotation_file = os.path.join(voc_annotations_path, voc_annotation)
        image_id = os.path.splitext(voc_annotation)[0]
        voc_image_file = os.path.join(voc_images_path, f"{image_id}.png")
        yolo_label_file = os.path.join(yolo_labels_path, f"{image_id}.txt")
        yolo_image_file = os.path.join(yolo_images_path, f"{image_id}.png")

        convert_voc_to_yolo(voc_annotation_file, yolo_label_file)
        if os.path.exists(voc_image_file):
            shutil.copy(voc_image_file, yolo_image_file)

print("VOC到YOLO格式转换完成！")

# 划分数据集
train_images_path = os.path.join(yolo_data_path, 'train', 'images')
train_labels_path = os.path.join(yolo_data_path, 'train', 'labels')
val_images_path = os.path.join(yolo_data_path, 'val', 'images')
val_labels_path = os.path.join(yolo_data_path, 'val', 'labels')
test_images_path = os.path.join(yolo_data_path, 'test', 'images')
test_labels_path = os.path.join(yolo_data_path, 'test', 'labels')

os.makedirs(train_images_path, exist_ok=True)
os.makedirs(train_labels_path, exist_ok=True)
os.makedirs(val_images_path, exist_ok=True)
os.makedirs(val_labels_path, exist_ok=True)
os.makedirs(test_images_path, exist_ok=True)
os.makedirs(test_labels_path, exist_ok=True)

# 获取所有图片文件名（不包含扩展名）
image_files = [f[:-4] for f in os.listdir(yolo_images_path) if f.endswith('.png')]

# 随机打乱文件顺序
random.shuffle(image_files)

# 划分数据集比例
train_ratio = 0.7
val_ratio = 0.2
test_ratio = 0.1

train_count = int(train_ratio * len(image_files))
val_count = int(val_ratio * len(image_files))
test_count = len(image_files) - train_count - val_count

train_files = image_files[:train_count]
val_files = image_files[train_count:train_count + val_count]
test_files = image_files[train_count + val_count:]

# 移动文件到相应的目录
def move_files(files, src_images_path, src_labels_path, dst_images_path, dst_labels_path):
    for file in tqdm(files):
        src_image_file = os.path.join(src_images_path, f"{file}.png")
        src_label_file = os.path.join(src_labels_path, f"{file}.txt")
        dst_image_file = os.path.join(dst_images_path, f"{file}.png")
        dst_label_file = os.path.join(dst_labels_path, f"{file}.txt")

        if os.path.exists(src_image_file) and os.path.exists(src_label_file):
            shutil.move(src_image_file, dst_image_file)
            shutil.move(src_label_file, dst_label_file)

# 移动训练集文件
print("移动训练集文件...")
move_files(train_files, yolo_images_path, yolo_labels_path, train_images_path, train_labels_path)
# 移动验证集文件
print("移动验证集文件...")
move_files(val_files, yolo_images_path, yolo_labels_path, val_images_path, val_labels_path)
# 移动测试集文件
print("移动测试集文件...")
move_files(test_files, yolo_images_path, yolo_labels_path, test_images_path, test_labels_path)

print("数据集划分完成！")

# 删除原始的 images 和 labels 文件夹
shutil.rmtree(yolo_images_path)
shutil.rmtree(yolo_labels_path)

print("原始 images 和 labels 文件夹删除完成！")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135