将图保存为numpy数组

Question

在Python和Matplotlib中，我们可以很简单地把图表显示成一个弹出窗口，或者把图表保存为PNG格式的文件。那么，怎么才能把图表保存成一个RGB格式的numpy数组呢？

Answer 1

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

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

Answer 2

对于@JUN_NETWORKS的回答，有一个更简单的选择。我们可以不把图像保存为png格式，而是用其他格式，比如raw或rgba，这样就可以省去cv2解码的步骤。

换句话说，实际的图像转换成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

这是一个很实用的小技巧，特别是在单元测试等场景中，当你需要对比一个保存的图像时，可以逐像素进行比较。

一种方法是使用 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.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))

将图保存为numpy数组

9 个回答

撰写回答