如何将 Matplotlib Axes 对象渲染到图像（作为 Numpy 数组）？

Question

有没有办法将特定 Axes 对象的内容呈现为图像，作为 Numpy 数组？我知道您可以对整个图形执行此操作，但我想获取特定轴的图像。

我尝试渲染的轴包含一个图像（使用 imshow 绘制），顶部绘制了一些线条。理想情况下，渲染的 ndarray 将只包含这些元素，没有刻度线、边框等。

更新：

我忘了提到我正在寻找一种不需要保存图像文件的解决方案。

我编写了以下示例，几乎可以实现这一点，只是它不保留图像分辨率。任何关于我可能做错的提示将不胜感激：

import matplotlib.pylab as plt
import numpy as np

def main():
  """Test for extracting pixel data from an Axes.

  This creates an image I, imshow()'s it to one Axes, then copies the pixel
  data out of that Axes to a numpy array I_copy, and imshow()'s the I_copy to
  another Axes.

  Problem: The two Axes *look* identical, but I does not equal I_copy.
  """

  fig, axes_pair = plt.subplots(1, 2)

  reds, greens = np.meshgrid(np.arange(0, 255), np.arange(0, 122))
  blues = np.zeros_like(reds)
  image = np.concatenate([x[..., np.newaxis] for x in (reds, greens, blues)],
                         axis=2)
  image = np.asarray(image, dtype=np.uint8)

  axes_pair[0].imshow(image)
  fig.canvas.draw()

  trans = axes_pair[0].figure.dpi_scale_trans.inverted()
  bbox = axes_pair[0].bbox.transformed(trans)
  bbox_contents = fig.canvas.copy_from_bbox(axes_pair[0].bbox)
  left, bottom, right, top = bbox_contents.get_extents()

  image_copy = np.fromstring(bbox_contents.to_string(),
                             dtype=np.uint8,
                             sep="")

  image_copy = image_copy.reshape([top - bottom, right - left, 4])
  image_copy = image_copy[..., :3]  # chop off alpha channel

  axes_pair[1].imshow(image_copy)

  print("Are the images perfectly equal? {}".format(np.all(image == image_copy)))

  plt.show()


if __name__ == '__main__':
  main()

Answer 1

一个想法可以是中间关闭轴，找出轴的边界框（以英寸为单位），然后使用bbox_inches 参数将图形保存到plt.savefig()。

如果需要一个 numpy 数组，可以使用plt.imread 再次读取保存的图像。

在此解决方案中，返回的数组的尺寸与在屏幕上绘制时轴的像素完全相同。

import matplotlib.pyplot as plt
import numpy as np; np.random.seed(0)

im = np.random.rand(16,16)

x = np.arange(9)
y = np.random.randint(1,14, size=(9,))
y2 = np.random.randint(1,7, size=(9,))

fig, (ax1, ax2) = plt.subplots(1,2)

ax1.imshow(im[:16,:9], cmap="viridis")
ax2.imshow(im[7:16,7:], cmap="YlGnBu")
ax1.plot(x, y, color="C3")
ax1.scatter(x, y, color="w")
ax2.plot(x, y2, color="C1")
ax2.scatter(x, y2, color="k")
ax1.set_xlabel("YlFasOcto")

def save_ax(ax, filename, **kwargs):
    ax.axis("off")
    ax.figure.canvas.draw()
    trans = ax.figure.dpi_scale_trans.inverted() 
    bbox = ax.bbox.transformed(trans)
    plt.savefig(filename, dpi="figure", bbox_inches=bbox,  **kwargs)
    ax.axis("on")
    im = plt.imread(filename)
    return im

arr = save_ax(ax1, __file__+".png")
print(arr.shape)
plt.show()

为了防止将文件保存到磁盘，可以使用 Stream 来保存数据。

import io

def save_ax_nosave(ax, **kwargs):
    ax.axis("off")
    ax.figure.canvas.draw()
    trans = ax.figure.dpi_scale_trans.inverted() 
    bbox = ax.bbox.transformed(trans)
    buff = io.BytesIO()
    plt.savefig(buff, format="png", dpi=ax.figure.dpi, bbox_inches=bbox,  **kwargs)
    ax.axis("on")
    buff.seek(0)
    im = plt.imread(buff )
    return im