MNIST
image visualization
data science
machine learning
Python

Show image from MNIST DataSet

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Showing an image from MNIST is usually the first sanity check before training a model. The dataset stores handwritten digits as 28 x 28 grayscale arrays, so the job is simply to load the data, pick one sample, and display it with a plotting library such as Matplotlib. The only subtle part is remembering that the data is numeric arrays, not already image files.

Load MNIST With Keras

A common way to access MNIST in Python is through tensorflow.keras.datasets.

python
1import matplotlib.pyplot as plt
2from tensorflow.keras.datasets import mnist
3
4(x_train, y_train), (x_test, y_test) = mnist.load_data()
5
6print(x_train.shape)  # (60000, 28, 28)
7print(y_train.shape)  # (60000,)

Each image is a 28 x 28 array of grayscale pixel intensities, and each label is the digit class from 0 through 9.

Display One Image

To show the first training image:

python
1import matplotlib.pyplot as plt
2from tensorflow.keras.datasets import mnist
3
4(x_train, y_train), _ = mnist.load_data()
5
6plt.imshow(x_train[0], cmap="gray")
7plt.title(f"Label: {y_train[0]}")
8plt.axis("off")
9plt.show()

cmap="gray" is important because MNIST images are grayscale. Without it, Matplotlib may apply a default color map that makes the digits look strange.

Display an Arbitrary Example by Index

You can inspect any sample by selecting its index.

python
1index = 123
2
3plt.imshow(x_train[index], cmap="gray")
4plt.title(f"Label: {y_train[index]}")
5plt.axis("off")
6plt.show()

This is useful when checking mislabeled data, examining model mistakes, or simply understanding dataset variation.

Show Multiple Digits in a Grid

For a better overview, display several images at once.

python
1import matplotlib.pyplot as plt
2from tensorflow.keras.datasets import mnist
3
4(x_train, y_train), _ = mnist.load_data()
5
6fig, axes = plt.subplots(2, 5, figsize=(10, 4))
7
8for i, ax in enumerate(axes.flat):
9    ax.imshow(x_train[i], cmap="gray")
10    ax.set_title(str(y_train[i]))
11    ax.axis("off")
12
13plt.tight_layout()
14plt.show()

A grid view is a quick way to verify that the dataset loaded correctly and that the labels line up with what you expect.

Understand the Shape Before Plotting

MNIST images are already two-dimensional arrays, so imshow can display them directly. That is different from many deep-learning pipelines where images are reshaped to include a channel dimension such as (28, 28, 1).

If your data has already been reshaped for a model, you may need to remove the extra dimension before plotting.

python
1image = x_train[0].reshape(28, 28, 1)
2plt.imshow(image.squeeze(), cmap="gray")
3plt.axis("off")
4plt.show()

That small detail avoids confusing shape errors.

Normalize for Training, Not for Display Logic

MNIST is often normalized before training:

python
x_train = x_train.astype("float32") / 255.0

Matplotlib can still display the image correctly, but it is helpful to remember the difference between visualization and model preprocessing.

  • raw MNIST pixel values range from 0 to 255
  • normalized values range from 0.0 to 1.0

Both can be plotted, but mixing the concepts can make it harder to reason about the pipeline.

Show Model Predictions Beside Images

Once you have a model, a useful extension is to display the prediction next to the ground-truth label.

python
1index = 7
2plt.imshow(x_test[index], cmap="gray")
3plt.title(f"True label: {y_test[index]}")
4plt.axis("off")
5plt.show()

After inference, you can update the title to include the predicted digit as well. This is one of the fastest ways to inspect classification mistakes manually.

Common Pitfalls

A common mistake is forgetting cmap="gray", which makes MNIST digits display with an unrelated color map.

Another issue is plotting reshaped tensors without removing the channel dimension first. If the array shape is not what imshow expects, the result may be an error or a misleading display.

Developers also sometimes normalize the images and then assume the plotted values should still look like 0 through 255. The image is fine, but the numeric range changed for training.

Finally, do not confuse the label array with the image array. y_train[index] is the digit class, while x_train[index] is the actual picture.

Summary

  • Load MNIST as numeric arrays, not image files.
  • Use plt.imshow(..., cmap="gray") to display a digit clearly.
  • Labels come from y_train or y_test, while images come from x_train or x_test.
  • Use grids to inspect multiple samples quickly.
  • Be mindful of array shape and normalization when moving between visualization and training code.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.