获取当前目录下所有的子目录
sub_dirs = [x[0] for x in os.walk(INPUT_DATA)]
# 得到的第一个目录是当前目录,不需要考虑
is_root_dir = True
for sub_dir in sub_dirs:
if is_root_dir:
is_root_dir = False
continue
# 获取当前目录下所有的有效图片文件。
extensions = ['jpg', 'jpeg', 'JPG', 'JPEG']
file_list = []
dir_name = os.path.basename(sub_dir)
for extension in extensions:
file_glob = os.path.join(INPUT_DATA, dir_name, '*.'+extension)
file_list.extend(glob.glob(file_glob))
if not file_list:
continue
# 通过目录名获取类别的名称。
label_name = dir_name.lower()
# 初始化当前类别的训练数据集、测试数据集和验证数据集
training_images = []
testing_images = []
validation_images = []
for file_name in file_list:
base_name = os.path.basename(file_name)
# 随机将数据分到训练数据集、测试数据集和验证数据集。
chance = np.random.randint(100)
if chance < validation_percentage:
validation_images.append(base_name)
elif chance < (testing_percentage + validation_percentage):
testing_images.append(base_name)
else:
training_images.append(base_name)
# 将当前类别的数据放入结果字典。
result[label_name] = {
'dir': dir_name,
'training': training_images,
'testing': testing_images,
'validation': validation_images
}