Transform numpy array to RGB image array

Question

Consider the following code:

import numpy as np
rand_matrix = np.random.rand(10,10)

which generates a 10x10 random matrix.

Following code to display as colour map:

import matplotlib.pyplot as plt
plt.imshow(rand_matrix)
plt.show()

I would like to get the RGB numpy array (no axis) from the object obtained from plt.imshow

In other words, if I save the image generated from plt.show, I would like to get the 3D RGB numpy array obtained from:

import matplotlib.image as mpimg
img=mpimg.imread('rand_matrix.png')

But without the need to save and load the image, which is computationally very expensive.

Thank you.

Answer 1

You can save time by saving to a io.BytesIO instead of to a file:

import io
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from PIL import Image

def ax_to_array(ax, **kwargs):
    fig = ax.figure
    frameon = ax.get_frame_on()
    ax.set_frame_on(False)
    with io.BytesIO() as memf:
        extent = ax.get_window_extent()
        extent = extent.transformed(fig.dpi_scale_trans.inverted())
        plt.axis('off')
        fig.savefig(memf, format='PNG', bbox_inches=extent, **kwargs)
        memf.seek(0)
        arr = mpimg.imread(memf)[::-1,...]
    ax.set_frame_on(frameon) 
    return arr.copy()

rand_matrix = np.random.rand(10,10)
fig, ax = plt.subplots()
ax.imshow(rand_matrix)
result = ax_to_array(ax)
# view using matplotlib
plt.show()
# view using PIL
result = (result * 255).astype('uint8')
img = Image.fromarray(result)
img.show()