深度学习常用的工具(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)