将图片和标签分成训练集和数据集同时 写入txt文档

import random
import os
#生成train.txt和val.txt
random.seed(2020)

xml_dir  = '/home/public/evening_yolo/data/train2017/labels'#标签文件地址
img_dir = '/home/public/evening_yolo/data/train2017/image'#图像文件地址

path_list = list()
for img in os.listdir(img_dir):
    img_path = os.path.join(img_dir,img)
    xml_path = os.path.join(xml_dir,img.replace('jpg', 'xml'))
    path_list.append((img_path, xml_path))
random.shuffle(path_list)
ratio = 0.9
train_f = open('/home/public/evening_yolo/data/train2017/train.txt','w') #生成训练文件
val_f = open('/home/public/evening_yolo/data/train2017/val.txt' ,'w')#生成验证文件

for i ,content in enumerate(path_list):
    img, xml = content
    text = img + ' ' + xml + '\n'
    if i < len(path_list) * ratio:
        train_f.write(text)
    else:
        val_f.write(text)
train_f.close()
val_f.close()

#生成标签文档
label = ['serial_num']#设置你想检测的类别
with open('/home/public/evening_yolo/data/train2017/label_list.txt', 'w') as f:
    for text in label:
        f.write(text+'\n')