加速matplotlib动画生成视频文件

Question

在Raspbian（树莓派2）上，以下是从我的脚本中剥离出来的最小示例，可以正确生成mp4文件：

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)

现在，文件的质量非常好，文件大小也很小，但是生成一个包含170帧的电影需要15分钟，这对我的应用程序来说是不可接受的。我正在寻找一种显著的加速方法，视频文件大小增加不是问题。

我认为视频制作中的瓶颈在于以png格式临时保存帧。在处理过程中，我可以看到png文件出现在我的工作目录中，CPU负载仅为25％。

请提供一个解决方案，该方案可能基于不同的软件包，而不仅仅是matplotlib.animation，例如OpenCV（已导入到我的项目中）或moviepy。

使用的版本：

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

- gaggio

Answer 1

Matplotlib 3.4 更新: 下面的方法可用于最新版本的 Matplotlib。然而，自本答案首次撰写以来，Matplotlib 的性能似乎有了重大提升，现在 Matplotlib 的 FFMpegWriter 速度与本方案的写入器相当。

原始答案：将动画保存到文件的瓶颈在于使用figure.savefig()。下面是一个自制的 matplotlib 的子类FFMpegWriter，灵感来自 gaggio 的答案。它不使用 savefig(因此忽略 savefig_kwargs)，但需要对您的动画脚本进行最小更改。

对于 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)) 

我能在半个时间内或更短的时间内创建动画，比使用默认的FFMpegWriter更快。 您可以按照此示例中所述使用它。

对于matplotlib >= 3.4 如果您将try块的最后一行更改为以下内容，则上面的代码将在matplotlib 3.4及以上版本中起作用：

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

即使用_proc.stdin而不是_frame_sink()。

- Aule Mahal

---

这个答案还是最新的吗？我尝试在当前的matplotlib版本中让它工作，但是animation.py自2018年以来明显已经改变了，例如_frame_sink()不再存在。您有什么建议吗，如何在今天实现它？ - CharlesT

我现在没有时间或设置来尝试这个，但是我在matplotlib代码中看到调用self._frame_sink()的部分被简单地替换为self._proc.stdin。你可以尝试一下这个方法。 编辑：如果有时间的话，我下周会尝试更新答案。 - Aule Mahal

Answer 2

一种显著改进的解决方案基于这篇文章的答案，大约可以将时间缩短10倍。

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)

我猜测，通过将原始图像数据传输到GStreamer，还可以类似地获得进一步的加速，在树莓派上现在可以从硬件编码中受益，详见此讨论。

- gaggio

---

这非常有帮助。写作类使用 fig.savefig() 保存到管道，速度非常慢。 - Wang

Answer 3

你应该能够使用其中一个编写器，直接流式传输到ffmpeg，但是其他一些问题正在发生。

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)

将此保存为test.py（这是您的代码的清理版本，我认为实际上无法运行，因为plt.plot返回一个Line2D对象列表，而列表没有plot方法）：

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

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

- tacaswell

---

谢谢您的纠正，实际上我的版本可以运行，因为使用了 pointplot[0]，但是您的版本更加简洁，我会相应地编辑原始版本。 - gaggio

无论如何，速度问题仍然存在... 你有没有注意到你的程序是否在工作目录中生成临时PNG文件？我假设你的版本也是在树莓派2上运行的，使用的是Raspbian发行版？也许我的原始问题表述不够清晰。 - gaggio

Answer 4

对于我的情况来说，花费的时间仍然太长，因此我使用了 @gaggio 的提议并使用 multiprocessing 进行了并行化。这至少在我的机器上有很大帮助，因为使用 ffmpeg 提出的解决方案似乎对于帧数没有线性时间复杂度。因此，我认为仅仅分块写入过程而不进行并行化也已经有所帮助。

假设您有一个 matplotlib 图形 fig，一个 animate(i) 函数用于更改动画的图形：

import multiprocessing
import math
import os

# divide into chunks (https://dev59.com/lnVC5IYBdhLWcg3wYQAp#312464)
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://dev59.com/FIzda4cB1Zd3GeqPjTRb#31315362)"""
    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)

我还没有实现清理临时文件的例程，但我猜这不是很难。

- sir_dance_a_lot