简体   繁体   中英

How to convert Matplotlib figure to PIL Image object (without saving image)

As the title states, I am trying to convert a fig to a PIL.Image . I am currently able to do so by first saving the fig to disk and then opening that file using Image.open() but the process is taking longer than expected and I am hoping that by skipping the saving locally step it will be a bit faster.

Here is what I have so far:

# build fig
figsize, dpi = self._calc_fig_size_res(img_height)
fig = plt.Figure(figsize=figsize)
canvas = FigureCanvas(fig)
ax = fig.add_subplot(111)
ax.imshow(torch.from_numpy(S).flip(0), cmap = cmap)
fig.subplots_adjust(left = 0, right = 1, bottom = 0, top = 1)
ax.axis('tight'); ax.axis('off')

# export
fig.savefig(export_path, dpi = dpi)

# open image as PIL object
img = Image.open(export_path)

I have tried doing this after I build the fig (it would be right before the export stage):

pil_img = Image.frombytes('RGB', canvas.get_width_height(), canvas.tostring_rgb())

But it's not showing the entire image. It looks like it's a crop of the top left corner, but it could just be a weird representation of the data -- I'm working with spectrograms so the images are fairly abstract.

EDIT # 2

PIL.Image.frombytes('RGB', 
fig.canvas.get_width_height(),fig.canvas.tostring_rgb())

takes around 2ms compared to the 35/40ms of the below.

This is the fastest way I can find so far.


I've been looking at this also today.

In the matplotlib docs the savefig function had this.

pil_kwargsdict, optional Additional keyword arguments that are passed to PIL.Image.save when saving the figure. Only applicable for formats that are saved using Pillow, ie JPEG, TIFF, and (if the keyword is set to a non-None value) PNG.

This must mean it's already a pil image before saving but I can't see it.

You could follow this

Matplotlib: save plot to numpy array

To get it into a numpy array and then do

PIL.Image.fromarray(array)

You might need to reverse the channels from BGR TO RGB with array [:, :, ::-1]

EDIT:

I've tested each way come up with so far.

import io
    
def save_plot_and_get():
    fig.savefig("test.jpg")
    img = cv2.imread("test.jpg")
    return PIL.Image.fromarray(img)
    
def buffer_plot_and_get():
    buf = io.BytesIO()
    fig.savefig(buf)
    buf.seek(0)
    return PIL.Image.open(buf)
    
def from_canvas():
    lst = list(fig.canvas.get_width_height())
    lst.append(3)
    return PIL.Image.fromarray(np.fromstring(fig.canvas.tostring_rgb(),dtype=np.uint8).reshape(lst))

Results

%timeit save_plot_and_get()

35.5 ms ± 148 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit save_plot_and_get()

35.5 ms ± 142 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit buffer_plot_and_get()

40.4 ms ± 152 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

I use the following function:

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

Example usage:

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

I flagged it as a duplicate and then closed it because I used the wrong link.

Anyway the answer may be here:

how to save a pylab figure into in-memory file which can be read into PIL image?

Unfortunately this did not lead to any speed increases, but I'll still post my specific solution below in case anyone has a similar issue:

# build fig
figsize, dpi = self._calc_fig_size_res(img_height)
fig = plt.Figure(figsize = figsize)
canvas = FigureCanvas(fig)
ax = fig.add_subplot(111)
ax.imshow(torch.from_numpy(S).flip(0), cmap = camp)
fig.subplots_adjust(left = 0, right = 1, bottom = 0, top = 1)
ax.axis('tight'); ax.axis('off')

# convert to PIL Image object
buf = io.BytesIO()
fig.savefig(buf, format='png', dpi = dpi)
buf.seek(0)
pil_img = deepcopy(Image.open(buf))
buf.close()

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM