I'm using the concept demonstrated in this paper. Their training data consists of "GOOD" images and "BAD" images. They train the AE using "BAD" images (X) to make it produce "GOOD" image. "BAD" images in this case maybe very similar to "GOOD" image but with small dent or scratch.

I've had success using ImageDataGenerator on mnist number data but in such case it is trained by using X -> X or if not using ImageDataGenerator the model fitting code will be something like model.fit(x_train, x_train).

However, in this AE problem, we want to train using X_good_and_defect -> X_good or model.fit(x_good_and_defect_train, x_good_train). Not sure how to achieve this using ImageDataGenerator.

I'm using keras's image data generator to load the images.

train_dir = r'chunks/training'
train_datagen = ImageDataGenerator(rescale=1 / 255)
train_generator = train_datagen.flow_from_directory(train_dir, target_size=(256,256), 
                   color_mode='grayscale', class_mode='input', batch_size=256)
                validation_data=(test_generator, train_generator),

As the following images from the paper shows, during training time you create patches from larger images. These patches have no defects and can therefore be seen as 'good' images. To get the accompanying 'bad' images with defects you synthetically generate these defects on the 'good' images.

enter image description here

Your data generator should therefore follow the following steps:

  1. Read in full size 'good' images
  2. Create patches from the full size images
  3. Synthetically add defects to the patches from step 2
  4. Return both the 'good' patches from step 2 and 'bad' patches from step 3

You then train your model using the 'bad' patches as the input and your 'good' patches as the output.

  • $\begingroup$ I understand, I'm having difficulty with the part "You then train your model using the 'bad' patches as the input and your 'good' patches as the output." using the train_generator though. Just not sure how. $\endgroup$
    – 40pro
    Feb 18 at 15:16
  • $\begingroup$ I see, I think you'd need to write a custom data generator that returns both the 'bad' and 'good' patches as ImageDataGenerator assumes that your data is an image and a label instead of an image and another image. I've personally found this website to be quite good. $\endgroup$
    – Oxbowerce
    Feb 18 at 15:21

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.