Test iterator and show images

test_dataset.py

import matplotlib.pyplot as plt 

init_vars = [tf.local_variables_initializer(), tf.global_variables_initializer()]

with tf.Session() as sess:
    sess.run([init_vars, iterator.initializer]) # Initialize variables and the iterator

    while 1:    # Iterate until we get out of range error!
        try:
            batch = sess.run(input_batch)
            print(batch.shape)  # Get batch dimensions
            plt.imshow(batch[0,:,:,0] , cmap='gray')
            plt.show()
        except tf.errors.OutOfRangeError:  # This exception is triggered when all batches are iterated
            print('All batches have been iterated!')
            break