如何将pylab图保存到内存文件中，以便可以读入PIL图像？

Question

以下是我从未成功的第一次尝试：

import cStringIO
import pylab
from PIL import Image
pylab.figure()
pylab.plot([1,2])
pylab.title("test")
buffer = cStringIO.StringIO()
pylab.savefig(buffer, format='png')
im = Image.open(buffer.read())
buffer.close()

错误信息显示：

Traceback (most recent call last):
  File "try.py", line 10, in <module>
    im = Image.open(buffer.read())
  File "/awesomepath/python2.7/site-packages/PIL/Image.py", line 1952, in open
    fp = __builtin__.open(fp, "rb")

有什么想法吗？我不希望解决方案涉及额外的包。

- nye17

Answer 1

记得调用buf.seek(0)，这样Image.open(buf)才会从buf的开头开始读取：

import io
from PIL import Image
import matplotlib.pyplot as plt

plt.figure()
plt.plot([1, 2])
plt.title("test")
buf = io.BytesIO()
plt.savefig(buf, format='png')
buf.seek(0)
im = Image.open(buf)
im.show()
buf.close()

- unutbu

---

太棒了！它像魔法一样运行！即使我用我的原始的 StringIO 替换 io.BytesIO 也是如此。你能提醒我为什么在这里选择使用前者吗？谢谢！ - nye17

对于Python2.6或更高版本，请使用io.BytesIO代替cStringIO.StringIO以实现向前兼容。在Python3中，cStringIO、StringIO模块已经被移除，它们的功能都在io模块中实现。 - unutbu

你如何为你的.png文件设置一个合理的文件名？它似乎只是称呼我的为fig.png - 谢谢。 - Little Bobby Tables

@josh: plt.savefig(buf, format='png') 将数据写入内存中的 BytesIO 对象，不会创建文件。通常使用 savefig 的方式是调用 plt.savefig('/path/to/anyfilenameyouwant.png')，这样可以指定文件名。 - unutbu

我正在从一个Flask应用程序的响应中发送此内容，以便Chrome等下载文件。也许有更好的方法可以选择我的文件名（下载的.png文件）。不过还是谢谢。 - Little Bobby Tables

如果你想将图像作为数组获取，可以不使用 PIL/pillow 库，而是使用 im = plt.imread(buf)。你也可以使用 plt.imsave('/path/to/saved/file.png', im) 将 im 数组保存到文件中。 - jmd_dk

Answer 2

我喜欢将其封装在一个函数中：

def fig2img(fig):
    """Convert a Matplotlib figure to a PIL Image and return it"""
    import io
    buf = io.BytesIO()
    fig.savefig(buf)
    buf.seek(0)
    img = Image.open(buf)
    return img

那么我可以很容易地这样调用它：

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

x = np.arange(-3,3)
plt.plot(x)
fig = plt.gcf()

img = fig2img(fig)
img.show()

- kotchwane

---

你也可以将图形保存为 fig.savefig(buf, bbox_inches='tight')，以去除默认情况下出现的厚白边，使得生成的图像更加紧凑和清晰，这是我个人的看法。 - Arturo Moncada-Torres

为什么你不关闭buf？ - TFAE