生成数据集train/val/test.txt文件列表

 生成数据集train/val/test.txt文件列表(前提是train/val/test数据图片是已经分好的)

import numpy as np

import os

 

base = './dataset/road/'

 

with open('dataset/road/labels.txt', 'w') as f:

    for i in range(20):

        f.write(str(i)+'\n')

 

imgs = os.listdir(base+'train_pic/')

np.random.seed(42)

np.random.shuffle(imgs)

val_num = int(0.1 * len(imgs))

 

with open(os.path.join('dataset/road/train_list.txt'), 'w') as f:

    for pt in imgs[:-val_num]:

        img = 'train_pic/'+pt

        ann = 'train_tag/'+pt.replace('.jpg', '.png'