Python Matplotlib:将打印保存到numpy数组

在Python和Matplotlib中，可以轻松地将绘图显示为弹出窗口或将绘图保存为PNG文件。如何将绘图保存为RGB格式的numpy数组？

这是一个方便的技巧，用于单元测试等，当您需要与保存的绘图进行像素到像素的比较时

一种方法是使用

fig.canvas.tostring_rgb

，然后使用合适的数据类型使用

numpy.fromstring

。还有其他方法，但这是我倾向于使用的方法

例如

---

有人提出这样一种方法

np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')

当然，这个代码是有效的。但是，输出的numpy阵列图像分辨率很低

我的提案代码是这样的

import io
import cv2
import numpy as np
import matplotlib.pyplot as plt

# plot sin wave
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)")


# define a function which returns an image as numpy array from figure
def get_img_from_fig(fig, dpi=180):
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=dpi)
    buf.seek(0)
    img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)
    buf.close()
    img = cv2.imdecode(img_arr, 1)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    return img

# you can get a high-resolution image as numpy array!!
plot_img_np = get_img_from_fig(fig)

此代码运行良好。

---

如果在dpi参数上设置一个大的数字，您可以获得一个高分辨率图像作为numpy数组。

对于@JUN_NETWORKS的答案，有一个更简单的选项。您可以使用其他格式，如

raw

或

rgba

，而不是将图形保存在

png

中，并跳过

cv2

解码步骤

换句话说，从绘图到numpy的实际转换归结为：

io_buf=io.BytesIO（）
图savefig（io_buf，格式='raw'，dpi=dpi）
io_buf.seek（0）
img_arr=np.reformate（np.frombuffer（io_buf.getvalue（），dtype=np.uint8），
newshape=（int（图bbox.bounds[3]），int（图bbox.bounds[2]），-1））
io_buf.close（）

---

希望，这会有所帮助。

如果有人想要一个即插即用的解决方案，而不需要修改任何先前的代码（获取pyplot figure和所有代码的引用），下面的内容对我很有用。只需在所有

pyplot

语句之后，即在

pyplot.show（）之前添加这个

是时候对您的解决方案进行基准测试了

导入io
导入matplotlib
matplotlib.use（'agg'）#关闭交互式后端
将matplotlib.pyplot作为plt导入
将numpy作为np导入
图，ax=plt.子批次（）
ax.绘图（范围（10））
def plot1（）：
图canvas.draw（）
data=np.frombuffer（图canvas.tostring_rgb（），dtype=np.uint8）
w、 h=图canvas.get_width_height（）
im=数据。重塑（（int（h），int（w），-1））
def plot2（）：
以io.BytesIO（）作为buff：
图savefig（buff，format='png'）
buff.seek（0）
im=plt.imread（buff）
def plot3（）：
以io.BytesIO（）作为buff：
图savefig（buff，format='raw'）
buff.seek（0）
data=np.frombuffer（buff.getvalue（），dtype=np.uint8）
w、 h=图canvas.get_width_height（）
im=数据。重塑（（int（h），int（w），-1））

在这种情况下，IO原始缓冲区是将matplotlib图形转换为numpy数组的最快速度

补充说明：


如果您无权访问图形，则始终可以从轴中提取图形：

fig=ax.图


如果需要
通道x高度x宽度格式的阵列，请执行以下操作

im=im.transpose（（2,0,1））


在Agg上工作，添加
matplotlib。在
导入matplotlib.pyplot作为plt之前使用（'Agg'）
以使用它。对于图像，画布增加了很大的边距，因此我发现在绘图之前插入
fig.tight_布局（pad=0）
很有用。对于具有线条和文本的图形，关闭抗锯齿功能也很重要。对于行
plt.setp（[ax.get\uxticklines（）+ax.get\uyticklines（）+ax.get\uxgridlines（）+ax.get\uygridlines（）]，antialiased=False）
和文本
mpl.rcParams['text.antialiased']=False
@JoeKington
np.fromstring
with
sep='
自版本1.14以来就不推荐使用。在将来的版本中，它应该替换为
data=np.frombuffer（图canvas.tostring_rgb（），dtype=np.uint8）
。如果遇到
“FigureCanvasGTKAgg”对象没有属性“renderer”
，请记住
matplotlib.use（'Agg'））
：我建议将导入语句与函数一起添加。@AnshulRai感谢您的宝贵建议！！我已经添加了有关导入、打印以及如何使用该函数的代码。我认为这个答案远远优于上面的答案：1）它生成高分辨率图像，2）不依赖cv2等外部软件包。我得到了一个整形错误“无法将3981312大小的数组整形为形状（480640，newaxis）”。有什么想法吗？事实上，这个答案正是我想要的！谢谢大家!@FabianHertwig-确保不仅像素数匹配，而且（颜色）通道数匹配。@FabianHertwig我也有同样的问题，下面是修复方法。创建fig时，必须将dpi设置为与保存时相同
fig=plt.figure（figsize=（16,4），dpi=128）
然后
fig.savefig（io_buf，format='raw'，dpi=128）
import io
import cv2
import numpy as np
import matplotlib.pyplot as plt

# plot sin wave
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)")


# define a function which returns an image as numpy array from figure
def get_img_from_fig(fig, dpi=180):
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=dpi)
    buf.seek(0)
    img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)
    buf.close()
    img = cv2.imdecode(img_arr, 1)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    return img

# you can get a high-resolution image as numpy array!!
plot_img_np = get_img_from_fig(fig)

canvas = pyplot.gca().figure.canvas
canvas.draw()
data = numpy.frombuffer(canvas.tostring_rgb(), dtype=numpy.uint8)
image = data.reshape(canvas.get_width_height()[::-1] + (3,))

>>> %timeit plot1()
34 ms ± 4.16 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
>>> %timeit plot2()
50.2 ms ± 234 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
>>> %timeit plot3()
16.4 ms ± 36 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)