How can I render a Matplotlib Axes object to an image (as a Numpy array)?

Question

Is there a way to render the contents of a particular Axes object to an image, as a Numpy array? I know you can do this with the entire figure, but I want to get the image of a

Answer 1

One idea can be to intermediately turn the axes off, find out the bounding box of the axes in inches and then save the figure using the bbox_inches argument to plt.savefig().

If a numpy array is wanted, one can then read in the saved image again using plt.imread.

In this solution the returned array has dimensions exactly as the axes has pixels when plotted on the screen.

import matplotlib.pyplot as plt
import numpy as np; np.random.seed(0)

im = np.random.rand(16,16)

x = np.arange(9)
y = np.random.randint(1,14, size=(9,))
y2 = np.random.randint(1,7, size=(9,))

fig, (ax1, ax2) = plt.subplots(1,2)

ax1.imshow(im[:16,:9], cmap="viridis")
ax2.imshow(im[7:16,7:], cmap="YlGnBu")
ax1.plot(x, y, color="C3")
ax1.scatter(x, y, color="w")
ax2.plot(x, y2, color="C1")
ax2.scatter(x, y2, color="k")
ax1.set_xlabel("YlFasOcto")

def save_ax(ax, filename, **kwargs):
    ax.axis("off")
    ax.figure.canvas.draw()
    trans = ax.figure.dpi_scale_trans.inverted() 
    bbox = ax.bbox.transformed(trans)
    plt.savefig(filename, dpi="figure", bbox_inches=bbox,  **kwargs)
    ax.axis("on")
    im = plt.imread(filename)
    return im

arr = save_ax(ax1, __file__+".png")
print(arr.shape)
plt.show()

In order to prevent saving a file to disk, one could use a Stream to save the data.

import io

def save_ax_nosave(ax, **kwargs):
    ax.axis("off")
    ax.figure.canvas.draw()
    trans = ax.figure.dpi_scale_trans.inverted() 
    bbox = ax.bbox.transformed(trans)
    buff = io.BytesIO()
    plt.savefig(buff, format="png", dpi=ax.figure.dpi, bbox_inches=bbox,  **kwargs)
    ax.axis("on")
    buff.seek(0)
    im = plt.imread(buff )
    return im