Matplotlib: 将绘图保存为 numpy 数组

在 Python 和 Matplotlib 中，很容易将情节显示为一个弹出窗口，或者将情节保存为一个 PNG 文件。如何才能保存的情节，而不是一个数字数组的 RGB 格式？

84966 次浏览
小开

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

一种方法是使用 fig.canvas.tostring_rgb，然后使用 numpy.fromstring和适当的 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.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
小开

有人提出这样一种方法

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

当然，这个代码可以工作。但是，输出的数组图像分辨率很低。

我的求婚密码是这个。

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 参数上设置一个较大的数字，就可以得到一个数字数组形式的高分辨率图像。

小开

对于@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()

霍普，这有帮助。

小开

如果有人想要一个即插即用的解决方案，而不需要修改任何先前的代码(获得对 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,))
小开

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

import io
import matplotlib
matplotlib.use('agg')  # turn off interactive backend
import matplotlib.pyplot as plt
import numpy as np


fig, ax = plt.subplots()
ax.plot(range(10))




def plot1():
fig.canvas.draw()
data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
w, h = fig.canvas.get_width_height()
im = data.reshape((int(h), int(w), -1))




def plot2():
with io.BytesIO() as buff:
fig.savefig(buff, format='png')
buff.seek(0)
im = plt.imread(buff)




def plot3():
with io.BytesIO() as buff:
fig.savefig(buff, format='raw')
buff.seek(0)
data = np.frombuffer(buff.getvalue(), dtype=np.uint8)
w, h = fig.canvas.get_width_height()
im = data.reshape((int(h), int(w), -1))

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

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

补充说明:

  • 如果你无法得到这个数字，你可以从轴上提取它:

    fig = ax.figure

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

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

小开

MoviePy 使得将图形转换为数字数组变得非常简单。它有一个称为 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
小开

正如乔 · 金顿所指出的，一种方法是在画布上绘画，将画布转换成字节串，然后将其重新塑造成正确的形状。

import matplotlib.pyplot as plt
import numpy as np
import math


plt.switch_backend('Agg')




def canvas2rgb_array(canvas):
"""Adapted from: https://stackoverflow.com/a/21940031/959926"""
canvas.draw()
buf = np.frombuffer(canvas.tostring_rgb(), dtype=np.uint8)
ncols, nrows = canvas.get_width_height()
scale = round(math.sqrt(buf.size / 3 / nrows / ncols))
return buf.reshape(scale * nrows, scale * ncols, 3)




# Make a simple plot to test with
t = np.arange(0.0, 2.0, 0.01)
s = 1 + np.sin(2 * np.pi * t)
fig, ax = plt.subplots()
ax.plot(t, s)


# Extract the plot as an array
plt_array = canvas2rgb_array(fig.canvas)
print(plt_array.shape)

但是，当 canvas.get_width_height()返回显示坐标中的宽度和高度时，有时会出现缩放问题，这个问题在这个答案中得到了解决。


提交答案