Practice Now
MNIST
Data Visualization
OpenCV
Matplotlib
Pyplot

Visualize MNIST dataset using OpenCV or Matplotlib/Pyplot

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
Computer VisionData Science

Introduction

The MNIST dataset contains 70,000 grayscale images of handwritten digits (28x28 pixels each). To visualize MNIST images, use matplotlib.pyplot.imshow() with cmap='gray' for quick plots, or cv2.imshow() with OpenCV for interactive windows. Matplotlib is better for grids, annotations, and notebook displays. OpenCV is better for real-time image processing pipelines. Both require the images to be reshaped from flat vectors (784,) to 2D arrays (28, 28) if loaded from certain sources.

Loading MNIST

python
1# Using TensorFlow/Keras
2from tensorflow.keras.datasets import mnist
3(X_train, y_train), (X_test, y_test) = mnist.load_data()
4print(X_train.shape)  # (60000, 28, 28)
5print(y_train.shape)  # (60000,)
6
7# Using PyTorch
8from torchvision import datasets
9train_data = datasets.MNIST(root='./data', train=True, download=True)
10image, label = train_data[0]  # PIL Image, int
11
12# Using scikit-learn (smaller subset — 1797 images, 8x8)
13from sklearn.datasets import load_digits
14digits = load_digits()
15print(digits.images.shape)  # (1797, 8, 8)

Visualize a Single Image with Matplotlib

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()

Visualize a Grid of Images

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

Show One Example Per Digit

python
1import matplotlib.pyplot as plt
2import numpy as np
3from tensorflow.keras.datasets import mnist
4
5(X_train, y_train), _ = mnist.load_data()
6
7fig, axes = plt.subplots(2, 5, figsize=(12, 5))
8for digit in range(10):
9    ax = axes[digit // 5, digit % 5]
10    idx = np.where(y_train == digit)[0][0]
11    ax.imshow(X_train[idx], cmap='gray')
12    ax.set_title(f"Digit {digit}")
13    ax.axis('off')
14plt.tight_layout()
15plt.show()

Visualize with OpenCV

python
1import cv2
2import numpy as np
3from tensorflow.keras.datasets import mnist
4
5(X_train, y_train), _ = mnist.load_data()
6
7# Single image — resize for visibility
8img = X_train[0]
9resized = cv2.resize(img, (280, 280), interpolation=cv2.INTER_NEAREST)
10cv2.imshow(f"Digit: {y_train[0]}", resized)
11cv2.waitKey(0)
12cv2.destroyAllWindows()
13
14# Grid of images using OpenCV
15def make_grid(images, labels, rows=4, cols=8, img_size=56):
16    grid = np.zeros((rows * img_size, cols * img_size), dtype=np.uint8)
17    for i in range(rows * cols):
18        r, c = divmod(i, cols)
19        resized = cv2.resize(images[i], (img_size, img_size),
20                             interpolation=cv2.INTER_NEAREST)
21        grid[r*img_size:(r+1)*img_size, c*img_size:(c+1)*img_size] = resized
22        # Add label text
23        cv2.putText(grid, str(labels[i]),
24                     (c*img_size + 2, r*img_size + 12),
25                     cv2.FONT_HERSHEY_SIMPLEX, 0.4, 255, 1)
26    return grid
27
28grid = make_grid(X_train, y_train)
29cv2.imshow("MNIST Grid", grid)
30cv2.waitKey(0)
31cv2.destroyAllWindows()

Visualize Pixel Value Distribution

python
1import matplotlib.pyplot as plt
2import numpy as np
3from tensorflow.keras.datasets import mnist
4
5(X_train, y_train), _ = mnist.load_data()
6
7# Histogram of pixel values for digit 0 vs digit 1
8fig, axes = plt.subplots(1, 2, figsize=(12, 4))
9
10digit_0 = X_train[y_train == 0].flatten()
11digit_1 = X_train[y_train == 1].flatten()
12
13axes[0].hist(digit_0[digit_0 > 0], bins=50, alpha=0.7)
14axes[0].set_title("Pixel Distribution: Digit 0")
15axes[0].set_xlabel("Pixel Value")
16
17axes[1].hist(digit_1[digit_1 > 0], bins=50, alpha=0.7, color='orange')
18axes[1].set_title("Pixel Distribution: Digit 1")
19axes[1].set_xlabel("Pixel Value")
20
21plt.tight_layout()
22plt.show()
23
24# Average image per digit class
25fig, axes = plt.subplots(2, 5, figsize=(12, 5))
26for digit in range(10):
27    ax = axes[digit // 5, digit % 5]
28    avg_img = X_train[y_train == digit].mean(axis=0)
29    ax.imshow(avg_img, cmap='hot')
30    ax.set_title(f"Avg: {digit}")
31    ax.axis('off')
32plt.suptitle("Average Image Per Digit")
33plt.tight_layout()
34plt.show()

Handling Flat Vectors

python
1import matplotlib.pyplot as plt
2import numpy as np
3
4# Some loaders return flat (784,) vectors instead of (28, 28) images
5flat_data = np.random.randint(0, 256, size=(100, 784))  # Simulated flat MNIST
6
7# Reshape to 28x28 for visualization
8img = flat_data[0].reshape(28, 28)
9plt.imshow(img, cmap='gray')
10plt.show()
11
12# Batch reshape
13images_2d = flat_data.reshape(-1, 28, 28)
14print(images_2d.shape)  # (100, 28, 28)

Common Pitfalls

  • Forgetting cmap='gray' in matplotlib: Without it, matplotlib uses a color map (viridis by default) for grayscale images, making digits appear in blue/yellow instead of white/black. Always pass cmap='gray' to imshow() for grayscale data.
  • Not reshaping flat vectors to 28x28: Some sources (scikit-learn's fetch_openml, CSV files) return MNIST as flat 784-element vectors. Passing a flat vector to imshow() fails with a shape error. Reshape with img.reshape(28, 28) before displaying.
  • Pixel values in wrong range for OpenCV: OpenCV imshow expects uint8 (0-255) for grayscale display. If your images are normalized to 0.0-1.0 (common after preprocessing), multiply by 255 and cast: (img * 255).astype(np.uint8). Otherwise the display appears all black or all white.
  • OpenCV window not appearing or freezing: cv2.imshow() requires cv2.waitKey() to process window events. Without it, the window appears frozen or never renders. Use cv2.waitKey(0) to wait for a keypress or cv2.waitKey(1) in a loop for real-time display.
  • Mixing up TensorFlow and PyTorch image formats: TensorFlow MNIST returns (batch, 28, 28) with values 0-255. PyTorch transforms often produce (batch, 1, 28, 28) with values 0.0-1.0. Before visualizing PyTorch tensors, squeeze the channel dimension and convert: img.squeeze().numpy().

Summary

  • Use plt.imshow(img, cmap='gray') for quick MNIST visualization in notebooks and scripts
  • Use cv2.imshow() with cv2.resize() for interactive or real-time display
  • Create grids with plt.subplots() to compare multiple digits at once
  • Reshape flat 784-element vectors to (28, 28) before displaying
  • Normalize pixel values to the correct range for each library (0-255 for OpenCV, 0-1 or 0-255 for matplotlib)
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
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.