Python - 将matplotlib图像转换为numpy.array 或 PIL.Image

1,939 阅读2分钟

这是我参与11月更文挑战的第9天,活动详情查看:2021最后一次更文挑战

matplotlib是python图像处理中让人又爱又恨的库。最近遇到了需要获取plt图像数据的需求,本文记录了将matplotlib图像转换为numpy.array 或 PIL.Image的方法。

众所周知,这个库处理图像会出现内存泄漏的问题,原想着将plt的图转出来用opencv存就好了,然而并没有,牢骚完毕。

转换思路

总体分为两步完成目标:

  • 将plt或fig对象转为argb string的对象
  • 将argb string对象图像转为array 或 Image

步骤一

区分对象为plt和fig的情况,具体使用哪种根据对象类型确定

转换plt对象为argb string编码对象

代码在plt对象中构建了图像内容,生成了plt图像,但还没有savefig 和 show:

例如:

 plt.figure()
 plt.imshow(img)
 # 引入 FigureCanvasAgg
 from matplotlib.backends.backend_agg import FigureCanvasAgg
 # 引入 Image
 import PIL.Image as Image
 # 将plt转化为numpy数据
 canvas = FigureCanvasAgg(plt.gcf())
 # 绘制图像
 canvas.draw()
 # 获取图像尺寸
 w, h = canvas.get_width_height()
 # 解码string 得到argb图像
 buf = np.fromstring(canvas.tostring_argb(), dtype=np.uint8)

转换fig对象为argb string编码对象

以 matplotlab 的 fig 对象为目标,获取 argb string编码图像

 # 引入 Image
 import PIL.Image as Image
 # 绘制图像
 fig.canvas.draw()
 # 获取图像尺寸
 w, h = fig.canvas.get_width_height()
 # 获取 argb 图像
 buf = np.fromstring(fig.canvas.tostring_argb(), dtype=np.uint8)

步骤二

转换argb string编码对象为PIL.Image或numpy.array图像

此时的argb string不是我们常见的uint8 w h rgb的图像,还需要进一步转化

 # 重构成w h 4(argb)图像
 buf.shape = (w, h, 4)
 # 转换为 RGBA
 buf = np.roll(buf, 3, axis=2)
 # 得到 Image RGBA图像对象 (需要Image对象的同学到此为止就可以了)
 image = Image.frombytes("RGBA", (w, h), buf.tostring())
 # 转换为numpy array rgba四通道数组
 image = np.asarray(image)
 # 转换为rgb图像
 rgb_image = image[:, :, :3]

参考资料