coco128数据集分割为训练集、验证集、测试集

coco128数据集已经转换为yolo格式（txt格式），将该数据集分割为train、valid、test三个子集。

• 解压出来的coco128里面，将images和labels下面的train2017里面的图片和标注文件分别移到images和labels里面
• 或者可以在代码中直接修改image_dir和label_dir

import os
import random
from tqdm import tqdm


# 指定 images 文件夹路径，images文件夹里面是图片，这里的路径换为自己的，一个是图片路径

image_dir = "../coco128/images"

# 指定 labels 文件夹路径，一个是标注文件路径

label_dir = "../coco128/labels"



# 创建一个空列表来存储有效图片的路径

valid_images = []

# 创建一个空列表来存储有效 label 的路径

valid_labels = []



# 遍历 images 文件夹下的所有图片

for image_name in os.listdir(image_dir):

    # 获取图片的完整路径

    image_path = os.path.join(image_dir, image_name)

    # 获取图片文件的扩展名

    ext = os.path.splitext(image_name)[-1]

    # 根据扩展名替换成对应的 label 文件名

    label_name = image_name.replace(ext, ".txt")

    # 获取对应 label 的完整路径

    label_path = os.path.join(label_dir, label_name)

    # 判断 label 是否存在

    if not os.path.exists(label_path):

        # 删除图片

        os.remove(image_path)

        print("deleted:", image_path)

    else:

        # 将图片路径添加到列表中

        valid_images.append(image_path)

        # 将label路径添加到列表中

        valid_labels.append(label_path)

        # print("valid:", image_path, label_path)



# 遍历每个有效图片路径

for i in tqdm(range(len(valid_images))):

    image_path = valid_images[i]

    label_path = valid_labels[i]

    # 随机生成一个概率

    r = random.random()

    # 判断图片应该移动到哪个文件夹

    # train：valid：test = 7:3:1

    if r < 0.1:

        # 移动到 test 文件夹，这里的三个路径换为自己的

        destination = "../datasets/test"

    elif r < 0.2:

        # 移动到 valid 文件夹

        destination = "../datasets/valid"

    else:

        # 移动到 train 文件夹

        destination = "../datasets/train"



    # 生成目标文件夹中图片的新路径

    image_destination_path = os.path.join(destination, "images", os.path.basename(image_path))

    # 移动图片到目标文件夹

    os.rename(image_path, image_destination_path)

    # 生成目标文件夹中 label 的新路径

    label_destination_path = os.path.join(destination, "labels", os.path.basename(label_path))

    # 移动 label 到目标文件夹

    os.rename(label_path, label_destination_path)



print("valid images:", valid_images)

#输出有效label路径列表

print("valid labels:", valid_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