Keras - 有没有办法管理由 ImageDataGenerator 的 flow_from_directory 函数生成的文件名?

Keras - Is there a way to manage the filenames generated by the flow_from_directory function of ImageDataGenerator?

由于标题是self-descriptive,我需要在数据增强后保留图像的原始文件名,这是由ImageDataGenerator flow_from_directory函数处理的 [= Keras 的 21=]。此要求背后的原因是文件名实际上代表标签,我将通过它们的名称将这些新图像移动到各自的文件夹中。请随时询问任何进一步的信息。

这是我的 ImageDataGenerator 以及我处理任务的方式:

aug = ImageDataGenerator(
    rotation_range=20,
    zoom_range=0.15,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.15,
    horizontal_flip=True,
    fill_mode="nearest")

i = 0
for batch in aug.flow_from_directory(extract_dir, batch_size=1, color_mode='grayscale', target_size=(28, 28),
                                     save_to_dir=extract_dir + '/augmented', save_prefix='aug'):
    i += 1
    if i == 100:
        break

不幸的是,没有一种简单的方法可以从 ImageDataGenerator.flow_from_directory 迭代器访问文件名。相反,您可以使用 ImageDataGenerator.flow 将增强应用于图像,然后使用另一个图像处理库手动保存增强图像,例如cv2.

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # suppress Tensorflow messages
from keras.preprocessing.image import ImageDataGenerator
import cv2
import numpy as np


image_folder = 'img_folder' # your image folder
print(f'my images: {os.listdir(image_folder)}')


aug_image_folder = 'aug_img_folder' # aug images will be saved here
os.makedirs(aug_image_folder,exist_ok=True)


# define your augmentations
aug = ImageDataGenerator(
    rotation_range=20,
    zoom_range=0.15,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.15,
    horizontal_flip=True,
    fill_mode="nearest")


# iterate over the image folder
image_names = os.listdir(image_folder)
for image_name in image_names:
    img = cv2.imread(f'{image_folder}/{image_name}',1)

    # get the augmented image
    aug_img_iterator = aug.flow(x=np.expand_dims(img,0),batch_size=1)
    aug_img=next(aug_img_iterator)

    # save the augmented image
    cv2.imwrite(f'{aug_image_folder}/{image_name}',aug_img[0,:,:,:])

print(f'aug images: {os.listdir(aug_image_folder)}')
my images: ['img_0.png', 'img_2.png', 'img_1.png']
aug images: ['img_0.png', 'img_2.png', 'img_1.png']