Yolov5-数据集的划分
import random
import os
import argparse
def get_opt():
parser = argparse.ArgumentParser()
parser.add_argument('--xml_path', default='D:/shuichi_test/VOC2007/Annotations/',
type=str, help='input xml file ')
parser.add_argument('--txt_path', default="D:/shuichi_test/VOC2007/ImageSets/Main/",
type=str, help='output txt file')
opt = parser.parse_args()
return opt
opt = get_opt()
xml_file = opt.xml_path
save_txt_file = opt.txt_path
if not os.path.exists(save_txt_file):
os.makedirs(save_txt_file)
total_xml = os.listdir(xml_file)
num = len(total_xml)
list_index = range(num)
train_val_percent = 1
train_percent = 0.99
tv = int(num * train_val_percent)
tr = int(tv * train_percent)
train_val = random.sample(list_index, tv)
train = random.sample(train_val, tr)
file_train_vale = open(save_txt_file + 'train_val.txt', 'w')
file_train = open(save_txt_file + "train.txt", 'w')
file_test = open(save_txt_file + "test.txt", 'w')
file_val = open(save_txt_file + "val.txt", 'w')
for i in list_index:
data_name = total_xml[i][:-4] + '\n'
if i in train_val:
file_train_vale.write(data_name)
if i in train:
file_train.write(data_name)
else:
file_val.write(data_name)
else:
file_test.write(data_name)
file_train_vale.close()
file_train.close()
file_test.close()
file_val.close()
效果显示
参考datawhale-yolov5训练到部署