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.

unutbu · Accepted Answer · 2016-08-03 10:24:31Z

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()

answered Aug 3, 2016 at 10:24

unutbu

886k197 gold badges1.9k silver badges1.7k bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

user3025898 Over a year ago

Thank you for your answer. This is working, but I wonder if it could be done with with open CV more efficiently?

Collectives™ on Stack Overflow

Transform numpy array to RGB image array

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related