将numpy数组转换为RGB图像数组

Question

请考虑以下代码：

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

生成10x10随机矩阵。

以下代码显示为彩色地图：

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

我想从plt.imshow获得的对象中获取RGB numpy数组（无轴）

换句话说，如果我保存从plt.show生成的图像，我想从下面获得3D RGB numpy数组：

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

但无需保存和加载图像，这在计算上非常昂贵。

谢谢。

Answer 1

您可以通过保存为io.BytesIO而不是文件来节省时间：

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