matplotlib 创建没有框架、轴的图形，使用颜色图绘制 2D 数组，将绘图保存到与输入大小相同的 numpy 数组

Question

我为此目的编写了一个函数：

1. 创建matplotlib图形，但不显示它
2. 没有框架、轴等
3. 使用用户传递的颜色图在图中绘制输入二维数组
4. 将颜色映射的二维数组从画布保存到numpy数组
5. 输出数组的大小应与输入数组的大小相同

对于类似于第 1-2 点或第 4 点的任务，有很多问题和答案；对我来说，自动化第 5 点也很重要。因此，我首先将 @joe-kington 的答案和 @matehat 的答案和评论中的部分结合起来，并进行了一些小的修改，我得到了：

def mk_cmapped_data(data, mpl_cmap_name):
    
    # This is to define figure & ouptput dimensions from input
    r, c = data.shape
    dpi = 72
    w = round(c/dpi, 2)
    h = round(r/dpi, 2)    
    
    # This part modified from @matehat's SO answer: 
    # /sf/answers/575322121/
    fig = plt.figure(frameon=False)
    fig.set_size_inches((w, h)) 
    ax = plt.Axes(fig, [0., 0., 1., 1.])
    ax.set_axis_off()
    fig.add_axes(ax)
    plt.set_cmap(mpl_cmap_name)
    ax.imshow(data, aspect='auto', cmap = mpl_cmap_name, interpolation = 'none')
    fig.canvas.draw()    
    
    # This part is to save the canvas to numpy array
    # Adapted rom Joe Kington's SO answer: 
    # /sf/answers/547534221/
    mat = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    mat = mat.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    mat = normalise(mat) # this is just using a helper function to normalize output range
    plt.close(fig=None)
    return mat

该函数完成了它应该做的事情并且足够快。我的问题是我是否可以以任何方式使其更加高效和/或更加Python化。

Answer 1

如果您希望 RGB 输出与输入数组的形状完全匹配，最简单的方法可能是不创建图形，而是直接使用颜色图对象。例如：

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

# Random data with a non 0-1 range.
data = 500 * np.random.random((100, 100)) - 200

# We'll use `LinearSegementedColormap` and `Normalize` instances directly
cmap = plt.get_cmap('viridis')
norm = plt.Normalize(data.min(), data.max())

# The norm instance scales data to a 0-1 range, cmap makes it RGB
rgb = cmap(norm(data))  

# MPL uses a 0-1 float RGB representation, so we'll scale to 0-255
rgb = (255 * rgb).astype(np.uint8) 

Image.fromarray(rgb).save('test.png')

请注意，您可能不希望执行将其另存为 PNG 的额外步骤，但我希望能够直观地显示结果。这正是一个 100x100 的图像，其中每个像素对应于原始输入数据。

这就是当您调用 imshow 时 matplotlib 在幕后所做的事情。数据首先运行一个Normalize实例，将其从原始范围缩放到 0-1。然后Colormap可以直接用0-1结果调用任意实例，将标量数据转换为RGB数据。