python 按比例分割数据集

用于yolov5训练。

import os
import random

images_path = '/data/my_set/images'
train_percent = 0.8

total_images = os.listdir(images_path)
random.shuffle(total_images) #打乱列表顺序

num = len(total_images)
carve_num = int(num*train_percent)
train = total_images[:carve_num]
val = total_images[carve_num:]

with open('train.txt', 'w') as f:
    for train_txt in train:
        f.write(os.path.join(images_path,train_txt) + '\n')
with open('val.txt', 'w') as f:
    for val_txt in val:
        f.write(os.path.join(images_path,val_txt) + '\n')