Keras: Get True labels (y_test) from ImageDataGenerator or predict_generator

Viewed 12553

I am using ImageDataGenerator().flow_from_directory(...) to generate batches of data from directories.

After the model builds successfully I'd like to get a two column array of True and Predicted class labels. With model.predict_generator(validation_generator, steps=NUM_STEPS) I can get a numpy array of predicted classes. Is it possible to have the predict_generator output the corresponding True class labels?

To add: validation_generator.classes does print the True labels but in the order that they are retrieved from the directory, it doesn't take into account the batching or sample expansion by augmentation.

3 Answers

You can get the prediction labels by:

 y_pred = numpy.rint(predictions)

and you can get the true labels by:

y_true = validation_generator.classes

You should set shuffle=False in the validation generator before this.

Finally, you can print confusion matrix by

print confusion_matrix(y_true, y_pred)

There's another, slightly "hackier" way, of retrieving the true labels as well. Note that this approach can handle when setting shuffle=True in your generator (it's generally speaking a good idea to shuffle your data - either if you do this manually where you've stored the data, or through the generator, which is probably easier). You will need your model for this approach though.

# Create lists for storing the predictions and labels
predictions = []
labels = []

# Get the total number of labels in generator 
# (i.e. the length of the dataset where the generator generates batches from)
n = len(generator.labels)

# Loop over the generator
for data, label in generator:
    # Make predictions on data using the model. Store the results.
    predictions.extend(model.predict(data).flatten())

    # Store corresponding labels
    labels.extend(label)

    # We have to break out from the generator when we've processed 
    # the entire once (otherwise we would end up with duplicates). 
    if (len(label) < generator.batch_size) and (len(predictions) == n):
        break

Your predictions and corresponding labels should now be stored in predictions and labels, respectively.

Lastly, remember that we should NOT add data augmentation on the validation and test sets/generators.

Using np.rint() method will get one hot coding result like [1., 0., 0.] which once I've tried creating a confusion matrix with confusion_matrix(y_true, y_pred) it caused error. Because validation_generator.classes returns class labels as a single number.

In order to get a class number for example 0, 1, 2 as class label specified, I have found the selected answer in this topic useful. here

Related