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 particular Axes.

The Axes I'm trying to render contains an image (drawn with imshow), with some lines plotted on top. Ideally, the rendered ndarray would contain just these elements, and no tick marks, borders, etc.

Update:

I forgot to mention that I'm looking for a solution that doesn't require saving an image file.

I've written the following example that almost achieves this, except it doesn't preserve image resolution. Any hints as to what I may be doing wrong would be greatly appreciated:

import matplotlib.pylab as plt
import numpy as np

def main():
  """Test for extracting pixel data from an Axes.

  This creates an image I, imshow()'s it to one Axes, then copies the pixel
  data out of that Axes to a numpy array I_copy, and imshow()'s the I_copy to
  another Axes.

  Problem: The two Axes *look* identical, but I does not equal I_copy.
  """

  fig, axes_pair = plt.subplots(1, 2)

  reds, greens = np.meshgrid(np.arange(0, 255), np.arange(0, 122))
  blues = np.zeros_like(reds)
  image = np.concatenate([x[..., np.newaxis] for x in (reds, greens, blues)],
                         axis=2)
  image = np.asarray(image, dtype=np.uint8)

  axes_pair[0].imshow(image)
  fig.canvas.draw()

  trans = axes_pair[0].figure.dpi_scale_trans.inverted()
  bbox = axes_pair[0].bbox.transformed(trans)
  bbox_contents = fig.canvas.copy_from_bbox(axes_pair[0].bbox)
  left, bottom, right, top = bbox_contents.get_extents()

  image_copy = np.fromstring(bbox_contents.to_string(),
                             dtype=np.uint8,
                             sep="")

  image_copy = image_copy.reshape([top - bottom, right - left, 4])
  image_copy = image_copy[..., :3]  # chop off alpha channel

  axes_pair[1].imshow(image_copy)

  print("Are the images perfectly equal? {}".format(np.all(image == image_copy)))

  plt.show()


if __name__ == '__main__':
  main()

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