#----------------------------------------#
#       生成猫狗数据集的样本列表文件
#----------------------------------------#
import os
import random


'''
    1.修改样本集目录 images_dir
    2.运行本文件
'''

#------------ Control Pannel ------------#
random.seed(0)

images_dir = 'E:/【AI】/Datasets/CV_ds/dogs_vs_cats/train/'

train_list_filepath = './model_data/train.txt'
valid_list_filepath = './model_data/valid.txt'
test_list_filepath  = './model_data/test.txt'

train_percent = 8/10
valid_percent = 1/10
test_percent  = 1/10
#----------------------------------------#


# 1.读取样本文件名
samples = os.listdir(images_dir)
random.shuffle(samples)

# 2.划分数据集
num_samples = len(samples)
s1 = int(train_percent * num_samples)
s2 = int(valid_percent * num_samples)
train_samples = samples[     :s1   ]
valid_samples = samples[s1   :s1+s2]
test_samples  = samples[s1+s2:     ]

# 3.生成样本列表文件
with open(train_list_filepath, 'w') as f:
    f.writelines([c+'\n' for c in train_samples])
with open(valid_list_filepath, 'w') as f:
    f.writelines([c+'\n' for c in valid_samples])
with open(test_list_filepath, 'w') as f:
    f.writelines([c+'\n' for c in test_samples])

# 4.输出数据集信息
print('Size of total, train, valid, test: {}, {}, {}, {}'.format(num_samples, s1, s2, num_samples-s2-s1))