Speedup matplotlib animation to video file

Question

On Raspbian (Raspberry Pi 2), the following minimal example stripped from my script correctly produces an mp4 file:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation

def anim_lift(x, y):

    #set up the figure
    fig = plt.figure(figsize=(15, 9))

    def animate(i):
        # update plot
        pointplot.set_data(x[i], y[i])

        return  pointplot

    # First frame
    ax0 = plt.plot(x,y)
    pointplot, = ax0.plot(x[0], y[0], 'or')

    anim = animation.FuncAnimation(fig, animate, repeat = False,
                                   frames=range(1,len(x)), 
                                   interval=200,
                                   blit=True, repeat_delay=1000)

    anim.save('out.mp4')
    plt.close(fig)

# Number of frames
nframes = 200

# Generate data
x = np.linspace(0, 100, num=nframes)
y = np.random.random_sample(np.size(x))

anim_lift(x, y)

Now, the file is produced with good quality and pretty small file size, but it takes 15 minutes to produce a 170 frames movie, which is not acceptable for my application. i'm looking for a significant speedup, video file size increase is not a problem.

I believe the bottleneck in the video production is in the temporary saving of the frames in png format. During processing I can see the png files apprearing in my working directory, with the CPU load at 25% only.

Please suggest a solution, that might also be based on a different package rather than simply matplotlib.animation, like OpenCV (which is anyway already imported in my project) or moviepy.

Versions in use:

• python 2.7.3
• matplotlib 1.1.1rc2
• ffmpeg 0.8.17-6:0.8.17-1+rpi1

Answer 1

You should be able to use one of the writers which will stream right to ffmpeg, but something else is going very wrong.

import matplotlib.pyplot as plt
from matplotlib import animation


def anim_lift(x, y):

    #set up the figure
    fig, ax = plt.subplots(figsize=(15, 9))

    def animate(i):
        # update plot
        pointplot.set_data(x[i], y[i])

        return [pointplot, ]

    # First frame
    pointplot, = ax.plot(x[0], y[0], 'or')
    ax.set_xlim([0, 200])
    ax.set_ylim([0, 200])
    anim = animation.FuncAnimation(fig, animate, repeat = False,
                                   frames=range(1,len(x)),
                                   interval=200,
                                   blit=True, repeat_delay=1000)

    anim.save('out.mp4')
    plt.close(fig)


x = list(range(170))
y = list(range(170))
anim_lift(x, y)

saving this as test.py (which is a cleaned up version of your code which I don't think actually runs because plt.plot returns a list of line2D objects and lists do not have a plot method) gives:

(dd_py3k) ✔ /tmp 
14:45 $ time python test.py

real    0m7.724s
user    0m9.887s
sys     0m0.547s

Answer 2

For my case it took still too long, wherefore I parallelized the proposal from @gaggio with multiprocessing. It helped for my machine at least by a magnitude since the proposed solutions with ffmpeg do not seem to have a linear time complexity for the number of frames. Therefore, I assume that chunking the writing process alone, without the parallelization, helps already.

Lets assume you have a matplotlib figure fig, an animate(i) function which changes the figure for the animation:

import multiprocessing
import math
import os

# divide into chunks (https://stackoverflow.com/a/312464/3253411)
def chunks(lst, n):
    """Yield successive n-sized chunks from lst."""
    for i in xrange(0, len(lst), n):
        yield lst[i:i + n]

# number of total frames
frames=1000
frame_iter=[i for i in range(frames)]

# distribute the frames over a set of equally sized chunks
chunk_size=math.ceil(number/multiprocessing.cpu_count())
frame_chunks=list(chunks(frames,chunk_size))

# get temporary video files to write to
filenames=["_temp_video_chunk_" + str(i) + ".mp4") for i in range(len(frame_chunks))]

def ani_to_mp4(frame_set, filename):
    """Animate figure fig for a defined frame set and save in filename (based n (https://stackoverflow.com/a/31315362/3253411)"""
    canvas_width, canvas_height = fig.canvas.get_width_height()

    # Open an ffmpeg process
    outf = os.path.join("results", filename)
    cmdstring = ('ffmpeg', 
                    '-y', '-r', '100', # fps
                    '-s', '%dx%d' % (canvas_width, canvas_height), # size of image string
                    '-pix_fmt', 'argb', # formats
                    '-f', 'rawvideo',  '-i', '-', # tell ffmpeg to expect raw video from the pipe
                    '-vcodec', 'mpeg4', outf) # output encoding
    p = subprocess.Popen(cmdstring, stdin=subprocess.PIPE)

    # Draw frames and write to the pipe
    for frame in frame_range:
        # draw the frame
        animate(frame)
        fig.canvas.draw()

        # extract the image as an ARGB string
        string = fig.canvas.tostring_argb()

        # write to pipe
        p.stdin.write(string)

    # Finish up
    p.communicate()

# take the function to write parallelized the animation chunks to the filenames
with multiprocessing.Pool() as pool:
    pool.starmap(ani_to_mp4, zip(frame_sets, filenames))

# write the filename list to a file
with open("filenames.txt", "w") as textfile:
    for filename in filenames:
        textfile.write("file '" + filename + "'\n")

# and use ffmpeg to concat the resulting mp4 files
cmdstring = ('ffmpeg', '-y',
                '-f', 'concat', 
                '-safe', '0', 
                '-i', "filenames.txt",
                '-c', 'copy', 
                'output.mp4') # output encoding
p = subprocess.Popen(cmdstring, stdin=subprocess.PIPE)

I haven't implemented a routine that cleans up the temporary files, but I guess that is not much work.

Answer 3

A much improved solution is based on the answers to this post reduces the time by a factor of 10 approximately.

import numpy as np
import matplotlib.pylab as plt
import matplotlib.animation as animation
import subprocess

def testSubprocess(x, y):

    #set up the figure
    fig = plt.figure(figsize=(15, 9))
    canvas_width, canvas_height = fig.canvas.get_width_height()

    # First frame
    ax0 = plt.plot(x,y)
    pointplot, = plt.plot(x[0], y[0], 'or')

    def update(frame):
        # your matplotlib code goes here
        pointplot.set_data(x[frame],y[frame])

    # Open an ffmpeg process
    outf = 'testSubprocess.mp4'
    cmdstring = ('ffmpeg', 
                 '-y', '-r', '1', # overwrite, 1fps
                 '-s', '%dx%d' % (canvas_width, canvas_height), # size of image string
                 '-pix_fmt', 'argb', # format
                 '-f', 'rawvideo',  '-i', '-', # tell ffmpeg to expect raw video from the pipe
                 '-vcodec', 'mpeg4', outf) # output encoding
    p = subprocess.Popen(cmdstring, stdin=subprocess.PIPE)

    # Draw frames and write to the pipe
    for frame in range(nframes):
        # draw the frame
        update(frame)
        fig.canvas.draw()

        # extract the image as an ARGB string
        string = fig.canvas.tostring_argb()

        # write to pipe
        p.stdin.write(string)

    # Finish up
    p.communicate()

# Number of frames
nframes = 200

# Generate data
x = np.linspace(0, 100, num=nframes)
y = np.random.random_sample(np.size(x))

testSubprocess(x, y)

I suspect further speedup might be obtained similarly by piping the raw image data to gstreamer which is now able to benefit from hardware encoding on the Raspberry Pi, see this discussion.

Answer 4

Matplotlib 3.4 update: The solution below can be adapted to work with the latest matplotlib versions. However, there seems to have been major performance improvements since this answer was first written and the speed of matplotlib's FFMpegWriter is now similar to this solution's writer.

Original answer: The bottleneck of saving an animation to file lies in the use of figure.savefig(). Here is a homemade subclass of matplotlib's FFMpegWriter, inspired by gaggio's answer. It doesn't use savefig (and thus ignores savefig_kwargs) but requires minimal changes to whatever your animation script are.

For matplotlib < 3.4

from matplotlib.animation import FFMpegWriter

class FasterFFMpegWriter(FFMpegWriter):
    '''FFMpeg-pipe writer bypassing figure.savefig.'''
    def __init__(self, **kwargs):
        '''Initialize the Writer object and sets the default frame_format.'''
        super().__init__(**kwargs)
        self.frame_format = 'argb'

    def grab_frame(self, **savefig_kwargs):
        '''Grab the image information from the figure and save as a movie frame.

        Doesn't use savefig to be faster: savefig_kwargs will be ignored.
        '''
        try:
            # re-adjust the figure size and dpi in case it has been changed by the
            # user.  We must ensure that every frame is the same size or
            # the movie will not save correctly.
            self.fig.set_size_inches(self._w, self._h)
            self.fig.set_dpi(self.dpi)
            # Draw and save the frame as an argb string to the pipe sink
            self.fig.canvas.draw()
            self._frame_sink().write(self.fig.canvas.tostring_argb()) 
        except (RuntimeError, IOError) as e:
            out, err = self._proc.communicate()
            raise IOError('Error saving animation to file (cause: {0}) '
                      'Stdout: {1} StdError: {2}. It may help to re-run '
                      'with --verbose-debug.'.format(e, out, err)) 

I was able to create animation in half the time or less than with the default FFMpegWriter.

You can use is as explained in this example.

For matplotlib >= 3.4

The code above will work with matplotlib 3.4 and above if you change the last line of the try block to:

self._proc.stdin.write(self.fig.canvas.tostring_argb())

i.e. using _proc.stdin instead of _frame_sink().