Save plot to numpy array

Question

In Python and Matplotlib, it is easy to either display the plot as a popup window or save the plot as a PNG file. How can I instead save the plot to a numpy array in RGB format?

Castrate

All the answers below don't offer the full solution - firstly one needs to prevent pyplot displaying when plotting it - which is the step before figure can be aqcuired – Honeymoon

Answer 1

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.

E.g.

import matplotlib.pyplot as plt
import numpy as np

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

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

# 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,))

Diary

Is this only supported on certain backend? Does not seem to be working with macosx backend (tostring_rgb) not found. – Diandre

Works on Agg, add matplotlib.use('agg') before import matplotlib.pyplot as plt to use it. – Diandre

With images, the canvas adds a big margin, so I found it useful to insert fig.tight_layout(pad=0) before drawing. – Geometrid

For figures with lines and text, it can also be important to turn antialiasing off. For lines plt.setp([ax.get_xticklines() + ax.get_yticklines() + ax.get_xgridlines() + ax.get_ygridlines()],antialiased=False) and for text mpl.rcParams['text.antialiased']=False – Adversaria

is it possible to store the the figure into to the same shape as the input dataset? – Dillon

There will be padding around figure. Various answers around for turning it off, but you can also just crop out the white from the resulting arrays with: data = data[:, np.where(~np.all(data==255, axis = 0))[0]] #cut all white rows data = data[np.where(~np.all(data==255, axis = 1,))[0], :] #and all white columns – Marutani

@JoeKington np.fromstring with sep='' is deprecated since version 1.14. It should be replaced with data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) in future versions – Esp

In case you run into 'FigureCanvasGTKAgg' object has no attribute 'renderer', remember to matplotlib.use('Agg'): https://mcmap.net/q/103358/-renderer-problems-using-matplotlib-from-within-a-script – Froude

instead matplotlib.use('agg') you can do plt.switch_backend('agg'), with this approach you won't have to run .use before import matplotlib.pyplot as plt – Threepence

I accidentally unvoted my upvote to @Threepence 's comment, but that is what I used in the end – Douville

Is print_to_buffer better than tostring_rgb? matplotlib.org/3.1.1/gallery/user_interfaces/canvasagg.html – Cockleboat

This does not work data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) Wrong shape – Krahmer

@MaratZakirov see my answer for a solution. – Outdoor

For some reason without the fig.canvas.draw() there are sizes mismatch, this was also mentioned by @Fabian Hertwig on Jonan Gueorguiev's answer – Phillada

Answer 2

There is a bit simpler option for @JUN_NETWORKS's answer. Instead of saving the figure in png, one can use other format, like raw or rgba and skip the cv2 decoding step.

In other words the actual plot-to-numpy conversion boils down to:

io_buf = io.BytesIO()
fig.savefig(io_buf, format='raw', dpi=DPI)
io_buf.seek(0)
img_arr = np.reshape(np.frombuffer(io_buf.getvalue(), dtype=np.uint8),
                     newshape=(int(fig.bbox.bounds[3]), int(fig.bbox.bounds[2]), -1))
io_buf.close()

Hope, this helps.

Seminar

I think this answer is far superior to the ones above: 1) It produces high-res images and 2) doesn't rely on external packages like cv2. – Kirk

I get a reshape error "cannot reshape array of size 3981312 into shape (480,640,newaxis)". Any ideas? – Christiechristin

Indeed this answer is exactly what I was looking for ! Thank you ! – Castile

@FabianHertwig - make sure that not only the number of pixels, but also the number of (color) channels match. – Seminar

@FabianHertwig I have the same problem, and here is the fix. When you create the fig you have to set dpi to be the same when you saved it fig = plt.figure(figsize=(16, 4), dpi=128) then fig.savefig(io_buf, format='raw', dpi=128) – Hilarity

How to read this in 3 channel i.e. rgb, something like (128, 128, 3). Right now it is coming as (128,128,4) ? – Adlay

This worked for me when omitting the dpi parameter, i.e. fig.savefig(io_buf, format='raw') – Nucleotidase

for anyone trying this method. It is very problematic (saving the image as 'raw' and to the io_buffer yields a shape that makes no sense). The top answer works far better – Minoan

Answer 3

Some people propose a method which is like this

np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')

Ofcourse, this code work. But, output numpy array image is so low resolution.

My proposal code is this.

import io
import cv2
import numpy as np
import matplotlib.pyplot as plt

# plot sin wave
fig = plt.figure()
ax = fig.add_subplot(111)

x = np.linspace(-np.pi, np.pi)

ax.set_xlim(-np.pi, np.pi)
ax.set_xlabel("x")
ax.set_ylabel("y")

ax.plot(x, np.sin(x), label="sin")

ax.legend()
ax.set_title("sin(x)")


# define a function which returns an image as numpy array from figure
def get_img_from_fig(fig, dpi=180):
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=dpi)
    buf.seek(0)
    img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)
    buf.close()
    img = cv2.imdecode(img_arr, 1)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    return img

# you can get a high-resolution image as numpy array!!
plot_img_np = get_img_from_fig(fig)

This code works well.
You can get a high-resolution image as a numpy array if you set a large number on the dpi argument.

Disclamation

I suggest adding the import statements along with the function. – Compression

@AnshulRai Thanks for your great advice!! I've added code about import, plot, and how to use the function. – Disclamation

Hey, I was wondering if there's a way I could get rid of x and y axis labels and just get the plot? – Very

Answer 4

Time to benchmark your solutions.

import io
import matplotlib
matplotlib.use('agg')  # turn off interactive backend
import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots()
ax.plot(range(10))


def plot1():
    fig.canvas.draw()
    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    w, h = fig.canvas.get_width_height()
    im = data.reshape((int(h), int(w), -1))


def plot2():
    with io.BytesIO() as buff:
        fig.savefig(buff, format='png')
        buff.seek(0)
        im = plt.imread(buff)


def plot3():
    with io.BytesIO() as buff:
        fig.savefig(buff, format='raw')
        buff.seek(0)
        data = np.frombuffer(buff.getvalue(), dtype=np.uint8)
    w, h = fig.canvas.get_width_height()
    im = data.reshape((int(h), int(w), -1))

>>> %timeit plot1()
34 ms ± 4.16 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
>>> %timeit plot2()
50.2 ms ± 234 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
>>> %timeit plot3()
16.4 ms ± 36 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Under this scenario, IO raw buffers are the fastest to convert a matplotlib figure to a numpy array.

Additional remarks:

• if you don't have an access to the figure, you can always extract it from the axes:

  fig = ax.figure

• if you need the array in the channel x height x width format, do

  im = im.transpose((2, 0, 1)).

Paries

Why do so many stackoverflowers obsess about benchmarks? – Outdoor

@Epimetheus, isn't it self-explanatory? the world only cares about speed and results' quality, a fast method can save companies. – Drive

Is this not missing a fig.canvas.draw() in the other two functions? – Indelible

@Indelible no, it's not needed in plot2 and plot3. Do you see an empty picture? – Paries

Answer 5

MoviePy makes converting a figure to a numpy array quite simple. It has a built-in function for this called mplfig_to_npimage(). You can use it like this:

from moviepy.video.io.bindings import mplfig_to_npimage
import matplotlib.pyplot as plt

fig = plt.figure()  # make a figure
numpy_fig = mplfig_to_npimage(fig)  # convert it to a numpy array

Durstin

As the shortest working solution Daniel Giger should get more upvotes. – Gagliano

But it requires installing another package :( – Bioenergetics

Answer 6

In case somebody wants a plug and play solution, without modifying any prior code (getting the reference to pyplot figure and all), the below worked for me. Just add this after all pyplot statements i.e. just before pyplot.show()

canvas = pyplot.gca().figure.canvas
canvas.draw()
data = numpy.frombuffer(canvas.tostring_rgb(), dtype=numpy.uint8)
image = data.reshape(canvas.get_width_height()[::-1] + (3,))

Bioenergetics

Answer 7

As Joe Kington has pointed out, one way is to draw on the canvas, convert the canvas to a byte string and then reshape it into the correct shape.

import matplotlib.pyplot as plt
import numpy as np
import math

plt.switch_backend('Agg')


def canvas2rgb_array(canvas):
    """Adapted from: https://mcmap.net/q/103359/-matplotlib-render-into-buffer-access-pixel-data"""
    canvas.draw()
    buf = np.frombuffer(canvas.tostring_rgb(), dtype=np.uint8)
    ncols, nrows = canvas.get_width_height()
    scale = round(math.sqrt(buf.size / 3 / nrows / ncols))
    return buf.reshape(scale * nrows, scale * ncols, 3)


# Make a simple plot to test with
t = np.arange(0.0, 2.0, 0.01)
s = 1 + np.sin(2 * np.pi * t)
fig, ax = plt.subplots()
ax.plot(t, s)

# Extract the plot as an array
plt_array = canvas2rgb_array(fig.canvas)
print(plt_array.shape)

However as canvas.get_width_height() returns width and height in display coordinates, there are sometimes scaling issues that are resolved in this answer.

Outdoor

Answer 8

Cleaned up version of the answer by Jonan Gueorguiev:

with io.BytesIO() as io_buf:
  fig.savefig(io_buf, format='raw', dpi=dpi)
  image = np.frombuffer(io_buf.getvalue(), np.uint8).reshape(
      int(fig.bbox.bounds[3]), int(fig.bbox.bounds[2]), -1)

Zygosis

Answer 9

import numpy as np 
import cv2
import time
import justpyplot as jplt

xs, ys = [], []
while(cv2.waitKey(1) != 27):
    xt = time.perf_counter() - t0
    yx = np.sin(xt)
    xs.append(xt)
    ys.append(yx)
    
    frame = np.full((500,470,3), (255,255,255), dtype=np.uint8)
    
    vals = np.array(ys)

    plotted_in_array = jplt.just_plot(frame, vals,title="sin() from Clock")
    
    cv2.imshow('np array plot', plotted_in_array)

The issues with all the matplotlib approaches is that matplotlib can still render and display plot even if you do plt.ioff() or return the figure and even if you do succeed while it behaves differently on a different platform(because matplotlib delegates it to backend depending on os) - you get a performance hit for getting plotted numpy array. I measured all previosly suggested matplotlib approaches and it rakes in milliseconds, most often dozens, sometimes even more milliseconds.

I couldn't find a simple library that just does it, had to write the thing myself. A plot to numpy in fully vectorized numpy(not a single loop) for all the parts such as scatter, connected, axis, grid, including size of the points and thickness and it does it in microseconds

https://github.com/bedbad/justpyplot

Honeymoon