[Python] 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.

Continue reading...