Matplotlib：将图保存到numpy数组

Question

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

Answer 1

当您需要对保存的图进行像素间比较时，这对于单元测试等来说是一个方便的技巧。

一种方法是使用dtype 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

有人提出这样的方法

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=180)

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