この記事の目的, 背景

matplotlibではプロットや画像の表示が出来ます. そしてそれを画像ファイルとして出力することも出来ます.
しかし, tensorboardへの記録やその後の画像の加工などがしたい場合には画像ファイルでは無く numpy array 形式で欲しい時があります.

しかし, 調べて出てきたやり方では低画質になってしまうという欠点がありました.

なので, 高画質な状態で numpy array 形式に出力するやり方を考えたのでそれについて書きます.

ベースとなるプロット

今回は定番のsin波をベースにしたいと思います.

import numpy as np
import matplotlib.pyplot as plt

fig = plt.figure()
ax = fig.add_subplot(111)

x = np.linspace(-np.pi, np.pi)

ax.set_xlim(-np.pi, np.pi)
ax.set_xlabel("x")
ax.set_ylabel("y")

ax.plot(x, np.sin(x), label="sin")

ax.legend()
ax.set_title("sin(x)")

プロットは以下のようになります.

f:id:JUN_NETWORKS:20191101012127p:plain — sin figure

従来の方法

savefig()

matplotlib の関数として savefig() というものがあり, それを使うと 画像ファイルとして出力できます.

fig.savefig("sin.png")

高画質で素晴らしいのですが, 我々が今欲しいのは numpy array 形式なのでこれは却下です.

もちろん, 出力したファイルを読み込めば numpy array 形式で扱えますけどね. なんか違うんだよなぁ.

プロットを文字列にして, それを numpy で読み込む

私が検索した限り, この方法が多かったです.

fig.canvas.draw()

data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))

stackoverflow.com

このコードは上手く動作します!! figure を numpy array で入手できてるしいいんじゃない?

だがしかし, 私はやり方では満足出来ませんでした.

以下が上記の data変数を画像として出力したものです.

f:id:JUN_NETWORKS:20191101013914p:plain — data変数の numpy array を画像として出力したもの

画質悪くない??

今回のプロットが簡単なのであまり気にならないのですが, これが大量のデータのプロットなどになるとプロットが潰れて見るに耐えないものになります.

今回の手法

従来の2つの手法はまとめると以下のようになります

savefig() : 画質は良いが, 画像ファイルが出力されるため numpy array にするのに出力画像を読み込まないといけないのでめんどくさい.
np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) : figure をファイルI/O無しで numpy array として入手出来るのはいいが, 画質がいまいち.

そして, 今回の私の手法がこれです. 上記2つの手法を組み合わせたとも言えなくないかも?

import io
import cv2

buf = io.BytesIO()  # インメモリのバイナリストリームを作成
fig.savefig(buf, format="png", dpi=180)  # matplotlibから出力される画像のバイナリデータをメモリに格納する.
buf.seek(0)  # ストリーム位置を先頭に戻る
img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)  # メモリからバイナリデータを読み込み, numpy array 形式に変換
buf.close()  # ストリームを閉じる(flushする)
img = cv2.imdecode(img_arr, 1)  # 画像のバイナリデータを復元する
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # cv2.imread() はBGR形式で読み込むのでRGBにする.