How to get combined image from matplotlib as variable?

Question

I have an image and heatmap. I plot both of them with im_show

axes_img1 = plt.imshow(255 * normalized_heat_map, alpha=alpha, cmap=cmap)
axes_img2 = plt.imshow(image, interpolation='nearest')
plt.show()

Both axes_img1 and axes_img2 are only for image or heatmap, plt.show() displays image with heatmap, in the way I want.

How can I get image (typically matplotlib outputs AxesImage type) that represents what plt.show() displays without saving any files?

Answer 1

plt.show() actually calls canvas.draw() and them shows the drawn (to buffer) image on the ui. hence, You need to call canvas.draw() and grab the buffer in numpy array.

def grab_buffer(fig):
    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    return data

fig,ax = plt.subplots(2,1) # 2 rows, 1 column
ax.flat[0].imshow(255 * normalized_heat_map, alpha=alpha, cmap=cmap)
ax.flat[1].imshow(image, interpolation='nearest')
fig.canvas.draw()
img_as_numpy_array = grab_buffer(fig)