[Python] moviepy + matplotlib で動画生成

これはなに

matplotlibで描画した画像をつなぎ合わせてmp4ファイルを作ります。

準備

pip install matplotlib moviepy

使用例

まずは公式ドキュメントにならって、小さい例で基本的な使い方を示します。

import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from moviepy.editor import VideoClip
from moviepy.video.io.bindings import mplfig_to_npimage

def make_frame(t):
    # 座標(t,t)に長方形を描画
    fig, ax = plt.subplots(figsize=(6,6))
    ax.add_patch(Rectangle((t, t), 1,2))
    ax.set_xlim(0,10)
    ax.set_ylim(0,10)
    # bitmapに変換. bmp.shape==(432,432,3)
    bmp = mplfig_to_npimage(fig)
    plt.close()
    return bmp
    
fname = "output/out2.mp4"
# 10秒, 30fps の計300フレームからなる動画を生成
animation = VideoClip(make_frame, duration = 10)
animation.write_videofile(fname, fps = 30)

インポートすべきいものが多いのでゴツいですが、要するにVideoClipオブジェクトに画像を描画するmake_frame()関数を渡して、動画を生成しているだけです。
make_frame()関数は1フレームの描画のたびに呼び出され、引数には秒数tが与えられます。返り値は各ピクセルのRGBの値を表す3次元配列を必要がありますが、これはmplfig_to_npimage関数を使うことでmatplotlibのfigureオブジェクトから生成することができます。

出力された動画は以下のようなものになります:

実用的な使い方

前節の使い方そのままだと、実際に大きなプログラムの中で描画～動画生成を行いたい場合にmake_frame関数の部分にすべての処理を記述することになってしまい、コードが複雑になってしまいます。あるいは、既にある大きなコードに動画出力の機能を足したい場合に、make_frame関数の中に既存のコードを移植することになってしまいます。

そこで実用的にはマルチスレッドにして

• メイン処理を行い、figを作るスレッド
• 作られたfigから動画を生成するスレッド

の2つに分けるのが見通しがよいです。

import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from moviepy.editor import VideoClip
from moviepy.video.io.bindings import mplfig_to_npimage
import threading
from queue import Queue

def func():
    """動画を生成
    """
    fname = "./out3.mp4"
    animation = VideoClip(make_frame, duration = 10)
    animation.write_videofile(fname, fps = 30)
def make_frame(t):
    """figqにbitmapが入ってきたらそれを取り出して返す
    """
    return figq.get()
def main():
    """メインの処理 + figqへのbitmapの追加
    """
    for i in range(310):
        fig, ax = plt.subplots(figsize=(6,6))
        ax.add_patch(Rectangle((i/30, i/30), 1,2))
        ax.set_xlim(0,10)
        ax.set_ylim(0,10)
        bmp = mplfig_to_npimage(fig)
        figq.put(bmp)
        plt.close()
    
figq = Queue()
thread = threading.Thread(target = func)
thread.start()
main()
thread.join()

先ほどとの変更点は、

• スレッドセーフなQueue(figq)を用意
• 動画の生成部分は別スレッド(thread)に切り出し
• make_frame()関数はfigqに入ってきたbitmapを取り出すだけの簡単なお仕事 (tは引数として与えられるが使わない)
• メイン処理～bitmapの生成・Queueへの追加をmain()関数に切り出し

といったあたりになります。結果としてはまったく同じ動画が出力されます。
既存のプログラムに動画生成の機能だけ付け加えたければ、main関数の

figq.put(bmp)

のあたりだけ追加すれば、あとの部分はコピペでOKだと思います。

実用例

AHC001のビジュアライザを作りました。

提出ではpypyで通るようにTEST=0としていますが、TEST=1とすると手元で動画を生成しながら動作します。
別途入力ファイルの準備も必要で、動かすのは面倒だと思うので眺める用にどうぞ。

できた動画は以下のような感じ:

そもそもアルゴリズムの性能がイマイチだったり、あくまで本記事のために書いたのでこのコンテストで必要な情報(広告の必須枠の位置とか)が載ってませんが。。

おまけ

jupyter notebookからであれば

import IPython
IPython.display.Video("./out3.mp4")

のようにしてセルの下で動画を見ることができて便利です。

注意点

• 必要なフレーム枚数をfigqにputしないと動画生成のスレッドが完了しないためthread.join()でハングアップしてしまいます。

  • また、2番目の例では手元の試行だとなぜか(10秒*30fps=300枚ではなく)301枚読み込もうとするらしく、301枚以上にする必要がありました。make_frame()が呼ばれるtの値を確認したところ、t=0が2回呼ばれているようです。いずれにしても、bitmapは少し多めにputしておくのが無難かと思います。

• 大量のフレームを投げつける場合メモリ周りで注意が必要です。

  • figqに要素が溜まりすぎるとメモリ消費が激しくなるので、そのような場合はpushする側がsleepなりして調整しましょう。
  • デフォルトのmatplotlibの設定だと画像を出力しようとして(?)figureオブジェクトのメモリが解放されず(?) Fail to allocate bitmapで落ちることがあるようです。この場合は例えば下記のようにコードの先頭に書くことで回避できます(できた):

import matplotlib
matplotlib.use("Agg")

(こちらのmatplotlibのissueを参考)

おわり

matplotlibだけでgif作る方が楽かなと思ってそっちも調べたのですが、汎用的な書き方があまりよくわからずでした。
mp4のほうがファイルサイズがずっと小さいのでmoviepy入れられるならこちらで良いのかなと。