How to save prediction result of the CNN model in the image format?

Question

I have been working with a CNN in Keras and I want to save the predictions as png images. This is my code to create the model and run the prediction:

 history = model.fit(
        train_generator,
        steps_per_epoch=train_steps,
        epochs=epochs,
        validation_data=validation_generator,
        validation_steps=valid_steps)

    pred= model.predict(X_test)

I used following code based on below link and the pred folder remains empty. ( How to save output (prediction) of the CNN model in the form of an image? )

images = []
for i in range(len(images)*2):
        pred.append(id_to_name[np.argmax(predictions[i // 2])])
        plt.imsave(('pred/ {}.png'.format(str(i // 2)+id_to_name[np.argmax(predictions[i // 2])])),images[i // 2])

Could you please tell how to save the model.predict predictions in the form of png images?

Answer 1

If your images have the format:

(batch_size, height, width, channels)

You can use tf.keras.preprocessing.image.save_img

import tensorflow as tf

images = tf.random.uniform((4, 224, 224, 3))

for i, image in enumerate(images, 1):
    tf.keras.preprocessing.image.save_img(f'my_picture_{i}.png', image)