Matplotlib:将绘图保存为numpy数组

Question

在Python和Matplotlib中,很容易将绘图显示为弹出窗口或将绘图保存为PNG文件.我怎样才能将绘图保存为RGB格式的numpy数组？

Answer 1

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

一种方法是使用fig.canvas.tostring_rgb,然后numpy.fromstring使用approriate dtype.还有其他方法,但这是我倾向于使用的方式.

例如

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 的答案有一个更简单的选项。除了将图形保存在 中png，还可以使用其他格式，例如raw或rgba并跳过cv2解码步骤。

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

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

\n

import io\nimport matplotlib\nmatplotlib.use(\'agg\')  # turn off interactive backend\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nfig, ax = plt.subplots()\nax.plot(range(10))\n\n\ndef plot1():\n    fig.canvas.draw()\n    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)\n    w, h = fig.canvas.get_width_height()\n    im = data.reshape((int(h), int(w), -1))\n\n\ndef plot2():\n    with io.BytesIO() as buff:\n        fig.savefig(buff, format=\'png\')\n        buff.seek(0)\n        im = plt.imread(buff)\n\n\ndef plot3():\n    with io.BytesIO() as buff:\n        fig.savefig(buff, format=\'raw\')\n        buff.seek(0)\n        data = np.frombuffer(buff.getvalue(), dtype=np.uint8)\n    w, h = fig.canvas.get_width_height()\n    im = data.reshape((int(h), int(w), -1))\n

\n

>>> %timeit plot1()\n34 ms \xc2\xb1 4.16 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 10 loops each)\n>>> %timeit plot2()\n50.2 ms \xc2\xb1 234 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 10 loops each)\n>>> %timeit plot3()\n16.4 ms \xc2\xb1 36 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 100 loops each)\n

\n

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

\n

附加说明：

\n

\n

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

  \n

  fig = ax.figure

  \n
\n

• 如果您需要格式的数组channel x height x width，请执行

  \n

  im = im.transpose((2, 0, 1))。

  \n
\n

\n

Answer 4

有些人提出了这样的方法

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 5

MoviePy使将图形转换为 numpy 数组变得非常简单。它有一个内置函数，称为mplfig_to_npimage(). 你可以这样使用它：

from moviepy.video.io.bindings import mplfig_to_npimage
import matplotlib.pyplot as plt

fig = plt.figure()  # make a figure
numpy_fig = mplfig_to_npimage(fig)  # convert it to a numpy array

Answer 6

如果有人想要一个即插即用的解决方案，而不修改任何先前的代码（获取对 pyplot 图和所有内容的引用），下面的内容对我有用。只需在所有语句之后添加此内容即可，pyplot即之前pyplot.show()

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,))