Pythonいぬ

pythonを使った画像処理に関する記事を書いていきます

matplotlibの画像を保存せずにnumpy arrayに変換

OpenCVなどでUSBカメラなどから得られた動画を流している際に、各フレームの解析結果をグラフにして動画と一緒にリアルタイムで表示したいときがある。このようなときにプロット図をいったんファイルに保存しながらやると面倒なので、できればメモリ上でのやりとりをしたい。

プロット画像をファイルに保存してから読み出す

まずは一度ファイルに保存してから再度読みだすスクリプト。この方法だと一時ファイルができる。

import numpy as np
import cv2

if __name__ == '__main__':

    x = np.arange(100)
    y = np.sin(np.pi*2*x/100)

    # 一度ファイルに保存                                                 
    plt.plot(x, y)
    plt.savefig("test.png") # ファイルに保存
    dst = cv2.imread('test.png') # ファイルから読み出し
    dst = dst[:,:,::-1] # BGR->RGB                                       
    plt.clf()

この方法でも遅くはならないけど、ファイルへの保存と読み込みをするので、保存先をどこにするかとか、実行中に同じ名前のファイルとかできたらどうするかが気になった。なので、Matplotlib でファイルを作らずにopencvにメモリを渡すだけで済む方法を考えてみた。

プロット画像をファイルに保存せずにnumpy arrayに変換する

メモリだけで渡すにはsavefigする際にファイルではなくバッファ(io.BytesIO)に入れて、numpyでuint8に変換してからopencvのcv2.imdecodeでデコードする。

import numpy as np
import cv2
import io

if __name__ == '__main__':

    x = np.arange(100)
    y = np.sin(np.pi*2*x/100)

    # プロット画像を直接メモリで渡す                                                   
    plt.plot(x, y)
    buf = io.BytesIO() # bufferを用意
    plt.savefig(buf, format='png') # bufferに保持
    enc = np.frombuffer(buf.getvalue(), dtype=np.uint8) # bufferからの読み出し
    dst = cv2.imdecode(enc, 1) # デコード
    dst = dst[:,:,::-1] # BGR->RGB
    plt.clf()

ファイル出力とやっていることはほとんど同じなんだけど、違いはバッファがファイルとして保存されないこと。