augmentation.py 文件源码

def augment(method, save_dir):
    train_filepaths = list(glob('./data/train/*/*.jpg'))
    labels = [path[13:-14] for path in train_filepaths]

    # Create label directories if they don't exist
    for label in labels:
        if not exists(save_dir + label):
            makedirs(save_dir + label)

    for file_path, label in tqdm(list(zip(train_filepaths, labels))):
        augmented_image = method(io.imread(file_path))
        io.imsave(save_dir + label + '/' + basename(file_path), augmented_image)