matplotlibの画像からPIL画像を取得する

matplotlibで描画した画像からPIL画像をファイルを経由せずに取得する方法について記載します。

matplotlibのグラフ画像の作成

とりあえず簡単なグラフを描画します。

test1.py

import os
import numpy as np
import matplotlib.pyplot as plt
import io
from PIL import Image

fig, ax = plt.subplots(1,1,figsize=(4,4))
x = np.linspace(-np.pi, np.pi)
ax.set_xlim(-np.pi, np.pi)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.grid()
ax.plot(x, np.cos(x), label="cos",color="lightgray")
ax.set_title("cos(X)")

# ファイルに保存
output_filename = "test1.png"
fig.savefig(output_filename)
plt.show() # jupyter notebookの場合
plt.close(fig)

PIL画像の取得

作成したグラフ画像をbuf = io.BytesIO()に出力します。
その後buf = io.BytesIO()からPIL画像を取得します。

test2.py

import os
import numpy as np
import matplotlib.pyplot as plt
import io
from PIL import Image

fig, ax = plt.subplots(1,1,figsize=(4,4))
x = np.linspace(-np.pi, np.pi)
ax.set_xlim(-np.pi, np.pi)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.grid()
ax.plot(x, np.cos(x), label="cos",color="lightgray")
ax.set_title("cos(X)")

# PIL画像に変換
buf = io.BytesIO()  
fig.savefig(buf, format="png", dpi=240) #解像度は調整
plt.close(fig)
buf.seek(0) # バッファの先頭に移動
img2=Image.open(buf).convert('RGB') # RGBAになるので変換

# PIL画像をファイルに保存
img2.save("test2.png") 
img2 # jupyter notebookの場合

注意点

ループを回す場合は、作成したグラフを必ずplt.close(fig)でクローズする必要があります。
メモリに蓄積されていくので注意が必要です。