Matplotlib: save plot to numpy array

This is a handy trick for unit tests and the like, when you need to do a pixel-to-pixel comparison with a saved plot.

One way is to use fig.canvas.tostring_rgb and then numpy.fromstring with the approriate dtype. There are other ways as well, but this is the one I tend to use.


import matplotlib.pyplot as plt
import numpy as np

# Make a random plot...
fig = plt.figure()

# If we haven't already shown or saved the plot, then we need to
# draw the figure first...

# Now we can save it to a numpy array.
data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))

Leave a Comment