获取当前目录下的有效图片

获取当前目录下所有的子目录

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
        }

猜你喜欢

转载自blog.csdn.net/weixin_41855385/article/details/84474861