Correctly reshaping MNIST images through keras

Question

If I download the CIFAR 10 Images in Keras via:

(x_train, y_train), (x_test, y_test) = cifar10.load_data()

# Getting shape
x_train.shape
>>> (50000, 32, 32, 3)

I can then plot each image by doing:

# Plot RGB image
plt.imshow(x_train[0])

#  Plot only one colour-channel e.g. R
plt.imshow(x_train[0][:,:,0])

Now I get the MNIST images via Keras by doing:

(X_train, y_train), (X_test, y_test) = mnist.load_data()

# getting shape
X_train.shape
>>> (60000, 28, 28)

However it doesn't have a depth channel which should be 1 as it's grayscale. If I reshape it using:

X_train = X_train.reshape(X_train.shape[0],28,28,1)

I can get the Neural Net to work but I can not plot it anymore. What is the correct way to reshape it so that I can still plot things?

Answer 1

As of v3.3 matplotlib's imshow now coerces 3d arrays of size MxNx1 into MxN for displaying, so you can plot your reshaped X_train without issue.