python - 将pyplot图形转换为数组

Question

我正在尝试将使用 pyplot 绘制的图形转换为数组，但我想在此之前消除绘图之外的任何空间。在我目前的方法中，我将图形保存到一个临时文件中（使用 的功能plt.savefig来消除绘图之外的任何空间，即使用bbox_inches='tight'and pad_inches = 0），然后从临时文件中加载图像。这是一个MWE：

from PIL import Image
import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots()
ax.plot([0,1], color='black', linewidth=4)
plt.xlim([0,1])
plt.ylim([0,1])
ax.set_aspect('equal', adjustable='box')
plt.axis('off')
plt.savefig('./tmp.png', bbox_inches='tight', pad_inches = 0)
plt.close()
img_size = 128
img = Image.open('./tmp.png')
X = np.array(img)

这种方法是不可取的，因为写入和读取文件需要时间。我知道以下直接从像素缓冲区到数组的方法：

from PIL import Image
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvas
import numpy as np

fig, ax = plt.subplots()
canvas = FigureCanvas(fig)
ax.plot([0,1], color='black', linewidth=4)
plt.xlim([0,1])
plt.ylim([0,1])
ax.set_aspect('equal', adjustable='box')
plt.axis('off')
canvas.draw()
X = np.array(canvas.renderer.buffer_rgba())

但是，使用这种方法，我不确定如何在转换为数组之前消除绘图周围的空间。是否有不涉及使用的等价bbox_inches='tight'物？pad_inches = 0plt.savefig()

标签： pythonnumpymatplotlibpython-imaging-library

Answer 1

改进的答案

这似乎适用于您的情况，应该很快。可能有更好的方法 - 如果有人知道更好的东西，我很乐意删除它：

#!/usr/bin/env python3

from PIL import Image
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvas
import numpy as np

fig, ax = plt.subplots()
canvas = FigureCanvas(fig)
ax.plot([0,1], color='red', linewidth=4)
plt.xlim([0,1])
plt.ylim([0,1])
ax.set_aspect('equal', adjustable='box')
plt.axis('off')
canvas.draw()
X = np.array(canvas.renderer.buffer_rgba())

上面的代码是你的，下面的代码是我的：

# Get width and height of cnvas for reshaping
w, h = canvas.get_width_height()
Y = np.frombuffer(X,dtype=np.uint8).reshape((h,w,4))[...,0:3]

# Work out extent of image by inverting and looking for black - ASSUMES CANVAS IS WHITE
extent = np.nonzero(~Y)

top    = extent[0].min()
bottom = extent[0].max()
left   = extent[1].min()
right  = extent[1].max()

tight_img = Y[top:bottom,left:right,:]

# Save as image just to test - you don't want this bit    
Image.fromarray(tight_img).save('tight.png')

原始答案

可能有更好的方法，但您可以通过写入基于内存的 BytesIO 来避免写入磁盘：

from io import BytesIO

buffer = BytesIO()
plt.savefig(buffer, format='png', bbox_inches='tight', pad_inches = 0)

然后做：

x = np.array(Image.open(buffer))

---

事实上，如果你使用：

plt.savefig(buffer, format='rgba', bbox_inches='tight', pad_inches = 0)

缓冲区已经有你的数组，你可以避免 PNG 编码/解码以及磁盘 I/O。唯一的问题是，因为它是原始的，所以我们不知道reshape()缓冲区中图像的尺寸。它实际上是在我的机器上，但我通过编写 PNG 并检查其宽度和高度来获得尺寸：

arr = buffer.getvalue()
x = np.frombuffer(arr, dtype=np.uint8).reshape((398,412,4))

如果有人想出更好的东西，我会删除它。