Save Plot To Numpy Array using Matplotlib

Saving a plot to a NumPy array in Python is a technique that bridges data visualization with array manipulation allowing for the direct storage of graphical plots as array representations, facilitating further computational analyses or modifications within a Python environment. Let’s learn how to Save Plot to NumPy Array using Matplotlib.

How to Save Plot to NumPy Array

To save a plot to a NumPy array, one must first create the plot using a plotting library like Matplotlib, then, utilizing `canvas.tostring_rgb()` method to capture the plot as an RGB string and reshape this data into a NumPy array with appropriate dimensions.

Essential steps include:

Plotting and capturing the plot’s RGB representation
Reshaping it into an array, ensuring seamless integration of visualization into data analysis workflows.

Method 1: Using fig.canvas.tostring_rgb and numpy.fromstring

In this approach, the plot is created and drawn on the canvas. The canvas is drawn to render the plot and the rendered canvas is then converted to a raw RGB buffer using ‘buf= fig.canvas.tostring_rgb()‘.
The raw buffer is converted into a NumPy array representing the image. Finally, it prints the shape of the image array and a portion of pixel values along with their RGB values.

Python3

import matplotlib.pyplot as plt

import numpy as np
 
fig, ax = plt.subplots()

ax.plot([1, 2, 3], [4, 5, 6])
fig.canvas.draw()
 
# Convert the canvas to a raw RGB buffer

buf = fig.canvas.tostring_rgb()

ncols, nrows = fig.canvas.get_width_height()

image = np.frombuffer(buf, dtype=np.uint8).reshape(nrows, ncols, 3)
 
print("Image shape:", image.shape)

print("First 3x3 pixels and RGB values:")

print(image[:3, :3, :])

Output:

Image shape: (480, 640, 3)
First 3x3 pixels and RGB values:
[[[255 255 255]
  [255 255 255]
  [255 255 255]]

 [[255 255 255]
  [255 255 255]
  [255 255 255]]

 [[255 255 255]
  [255 255 255]
  [255 255 255]]]

The dimensions of the image are (480, 640, 3), which indicates:
480: This is the height of the image in pixels.
640: This is the width of the image in pixels.
3: This represents the number of color channels in the image. Since the value is 3, it’s likely an RGB image where each pixel has three values representing red, green, and blue intensities.

Method 2: Saving the plot to a io.BytesIO object

This method involves creating a plot with Matplotlib and saving it to a BytesIO object in memory (as PNG) a temporary buffer in memory, instead of directly to a file.
The BytesIO object is then read into a PIL Image, using the Pillow (PIL Fork) library which is converted to a NumPy array. This approach is efficient for converting plots to arrays without needing to save and read from disk.

Python3

import matplotlib.pyplot as plt

import numpy as np

import io

from PIL import Image
 
fig, ax = plt.subplots()

ax.plot([1, 2, 3], [4, 5, 6])
 
# Create a bytes buffer to save the plot

buf = io.BytesIO()

plt.savefig(buf, format='png')

buf.seek(0)
 
# Open the PNG image from the buffer and convert it to a NumPy array

image = np.array(Image.open(buf))
# Close the buffer
buf.close()
 
print("Image shape:", image.shape)

print("First 3x3 pixels and RGB values:")

print(image[:3, :3, :])

Output:

Image shape: (480, 640, 4)
First 3x3 pixels and RGB values:
[[[255 255 255 255]
  [255 255 255 255]
  [255 255 255 255]]

 [[255 255 255 255]
  [255 255 255 255]
  [255 255 255 255]]

 [[255 255 255 255]
  [255 255 255 255]
  [255 255 255 255]]]

The output (480, 640, 4) gives a NumPy array representing an image with a height of 480 pixels, a width of 640 pixels, and four color channels (RGBA). While the alpha channel might not be explicitly used in the plot data itself, its presence could be due to Matplotlib’s handling of PNG transparency or the behavior of the image processing library used.

Matplotlib with PNG transparency: When saving a Matplotlib plot to PNG format using plt.savefig with transparency enabled, the resulting image might have an alpha channel even if the plot itself doesn’t explicitly use transparency. This behavior depends on Matplotlib’s internal handling of PNG images.
External library (PIL): While using an external library like PIL to open the image, it might add an alpha channel even if the original PNG from Matplotlib didn’t have it. Some image processing libraries like PIL might assume or add an alpha channel by default for consistency.

Conclusion

In conclusion, converting plots to NumPy arrays in Python enhances data visualization by integrating it with array-based computation, offering flexibility across various applications.

Article Tags :

AI-ML-DS

Python

Python-matplotlib