Matplotlib:将绘图保存到numpy数组

Question

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

Answer 1

当您需要对保存的绘图进行像素到像素的比较时，这对于单元测试之类的东西来说是一个很方便的技巧。

一种方法是先使用fig.canvas.tostring_rgb，然后使用带有适当数据类型的numpy.fromstring。还有其他方法，但这是我倾向于使用的方法。

例如。

import matplotlib.pyplot as plt
import numpy as np

# Make a random plot...
fig = plt.figure()
fig.add_subplot(111)

# If we haven't already shown or saved the plot, then we need to
# draw the figure first...
fig.canvas.draw()

# Now we can save it to a numpy array.
data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))

Answer 2

对于@JUN_NETWORKS的答案，有一个更简单的选择。可以使用其他格式，如raw或rgba，跳过cv2解码步骤，而不是将图形保存为png格式。

换句话说，实际的plot-to-numpy转换可以归结为：

io_buf = io.BytesIO()
fig.savefig(io_buf, format='raw', dpi=DPI)
io_buf.seek(0)
img_arr = np.reshape(np.frombuffer(io_buf.getvalue(), dtype=np.uint8),
                     newshape=(int(fig.bbox.bounds[3]), int(fig.bbox.bounds[2]), -1))
io_buf.close()

霍普，这有帮助。

Answer 3

有些人提出了这样一种方法

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数组。