Faster Grabbing of Matplotlib Animation Frames as Numpy Ndarrays

This is a continuation of the previous post on Quickly generating high quality scientific videos with Python

In the previous post we explored optimal parameters for video encoding in generating high resolution, decently compressed scientific animations at high speed. But we still need to generate these frames. Specifically, the question here is:

How to efficiently convert a matplotlib Figure to a Numpy RGB or RGBA ndarray?

For each frame we will be plotting one simple waveform. The figure will be generated at 720p resolution.

I am also asking for a Numpy ndarray as output as I might need to post-process or concatenate this image on top of other, so I don’t want it to be directly written to an MP4 file. (However if only a video file or even an interactive HTML preview of the animation is needed, check out matplotlib.animation which is also compatible with Jupyter Notebooks)

The reason I am sticking to matplotlib is because it is something that everyone is familar with (no learing curve) and can be reliably used in headless environments. There are supposely faster libraries out there, such as:

• pyqtgraph
• pyvis

mplfig_to_npimage from moviepy

MoviePy comes with a function named mplfig_to_npimage exactly for this purpose. The key step is generate a canvas and use canvas.tostring_rgb() to get an RGB array out. This is the rate limiting step in the whole process.

The example below is taken from moviepy documentation but is optimized by using line.set_data which is much faster than ax.clear() on every loop.

import matplotlib.pyplot as plt
import numpy as np
from moviepy.editor import VideoClip
from moviepy.video.io.bindings import mplfig_to_npimage

x = np.linspace(-2, 2, 200)

duration = 2

fig, ax = plt.subplots(figsize=(12.80, 7.20), dpi=100)
[line] = ax.plot(x, np.sin(x))
def make_frame(t):
    line.set_data(x, np.sin(x+t))
    return mplfig_to_npimage(fig)

%time regular_frames = [make_frame(i) for i in range(300)]
CPU times: user 12.6 s, sys: 702 ms, total: 13.3 s
Wall time: 13.3 s

So the performance is 12.6 seconds per 300 frames or 24 FPS.

savefig to buffer

Unsatisfied with the performance, I soon realized that the matplotlib animation code does not actually use the tostring_rgb approach but rather uses savefig to stdin to grab the frame. So below I switched out the implementation and used savefig to save the figure in memory and then read it by numpy.

from io import BytesIO

fig, ax = plt.subplots(figsize=(12.80, 7.20), dpi=100)
[line] = ax.plot(x, np.sin(x))

fig.canvas.draw()  # update/draw the elements
# get the width and the height to resize the matrix
l, b, w, h = fig.canvas.figure.bbox.bounds
w, h = int(w), int(h)

def make_frame_faster(t):
    line.set_data(x, np.sin(x+t))

    # Saving to buffer and then 
    # reading it back out to numpy
    bio = BytesIO()
    fig.savefig(bio, format="rgba")
    bio.seek(0)
    image = np.frombuffer(bio.read(), dtype=np.uint8)
    image = image.reshape(h, w, -1)
    return image

%time faster_frames = [make_frame_faster(i) for i in range(300)]
CPU times: user 8.33 s, sys: 674 ms, total: 9 s
Wall time: 8.89 s

This proved to be much faster at 8.33 seconds per 300 frames or 36 FPS. That’s 50% faster!

We can also see below on the left and right that the frames output by the two methods are indeed identical. (In fact the second method also preserves the fourth alpha channel)

fig, ax = plt.subplots(5, 2, figsize=(8.5, 11))
for i, frame_i in enumerate(range(0, 100, 20)):
    ax[i, 0].imshow(regular_frames[frame_i])
    ax[i, 1].imshow(faster_frames[frame_i])

Conclusion

So there is a faster way than mplfig_to_npimage, by directly savefig(..., format="rgba") to memory.

Of course right now this is a really simple example of a waveform. If you have something more complicated, you will want to learn about “Blitting” which is covered in matplotlib animation guide.

This test was also drafted on a headless Colab notebook. I am unsure if the performance would be affected by the matplotlib backend used.