深度学习常用的工具（1）根据比率来随机生成训练的train.txt和val.txt

深度学习常用的工具（1）
根据比率来随机划分生成训练的train.txt和val.txt

"""
根据比率来随机生成训练的train.txt和val.txt
"""
import os
import random


def gen_train_val(path, ratio):
    imgs = os.listdir(path)
    train_num = int(len(imgs) * ratio)
    train_imgs = random.sample(imgs, train_num)
    val_imgs = []
    for img in imgs:
        if img not in train_imgs:
            val_imgs.append(img)
    with open(r'D:\work\project\train.txt', 'a', encoding='utf-8') as f:
        for train_img in train_imgs:
            f.write(train_img.split('.')[0]+ '\n')
    with open(r'D:\work\project\val.txt', 'a', encoding='utf-8') as f:
        for val_img in val_imgs:
            f.write(val_img.split('.')[0]+'\n')


if __name__ == '__main__':
    # 标注的图片数据集路径
    path = r'D:\wx\img'
    # 比率
    ratio = 0.9
    gen_train_val(path, ratio)