变换numpy的阵列到RGB图像阵列

Question

考虑下面的代码：

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

其产生10×10随机矩阵。

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

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

我想获得从plt.imshow

获得换句话说，对象的RGB numpy的阵列（无轴），如果我保存从plt.show生成的图像，我想获得从所获得的三维RGB numpy的阵列：

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

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

谢谢。

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