如何自定义0-255灰度值转RGBA四通道Viridis热力图的映射函数

核心原理 #

所有matplotlib内置连续色图（包括viridis）本质都是分段线性插值查找表，没有特殊黑箱逻辑：

• 色图预先定义了一组锚点，每个锚点对应「归一化输入位置 + R/G/B/A四个通道的输出值」
• 映射时先把输入值归一化到[0,1]区间，找到值落在的两个相邻锚点区间，对四个通道做线性插值，就能得到对应颜色
• viridis官方内置了256个均匀采样的控制点，刚好匹配0-255的灰度值输入范围，不需要额外计算插值就能直接索引。

实现代码 #

首先注意你提供的测试代码缺失numpy导入，运行前需要补上import numpy as np。下面提供两种实现，第一种和官方cm.viridis输出100%一致，适合生产使用；第二种手动实现插值逻辑，方便你理解通用色图映射规则。

方法1：直接调用内置查找表（推荐） #

matplotlib的所有colormap对象都自带colors属性，是shape为(256,4)的数组，索引n直接对应输入灰度值n的RGBA输出，用numpy向量化索引可以实现毫秒级映射，不需要逐像素循环。

import numpy as np
import matplotlib.pyplot as plt
import requests
from PIL import Image
from io import BytesIO
from matplotlib import cm

# 预加载viridis的256色查找表，值范围为0-1的浮点数
viridis_lut = cm.viridis(np.arange(256))

def custom_viridis(gray_arr):
    """
    输入：shape为(H,W)、取值范围0-255的灰度数组
    输出：shape为(H,W,4)的RGBA热力图数组，和cm.viridis输出完全一致
    """
    return viridis_lut[gray_arr]

# 测试验证
img_src = 'https://unsplash.it/500/300'
response = requests.get(img_src)
imgarray = Image.open(BytesIO(response.content))
imgarray = np.asarray(imgarray.convert('L'))

official_output = cm.viridis(imgarray)
custom_output = custom_viridis(imgarray)
# 验证一致性，输出为True即结果完全匹配
print(np.allclose(official_output, custom_output))
plt.imshow(custom_output)
plt.show()

方法2：手动实现通用映射逻辑 #

如果需要完全脱离matplotlib封装、自定义色图规则，可以用通用的分段插值逻辑实现，这个逻辑对所有连续色图都生效：

1. 将0-255的输入灰度值归一化到[0,1]区间
2. 准备色图的锚点列表，记录每个锚点的归一化位置和对应R/G/B值
3. 对每个输入值，找到相邻的两个锚点，按距离权重线性插值计算三个通道值，Alpha通道默认设为1（完全不透明）

# viridis官方核心锚点（归一化位置, R, G, B），来自matplotlib源码色图定义
viridis_anchors = np.array([
    [0.0, 0.267004, 0.004874, 0.329415],
    [0.1, 0.282327, 0.140926, 0.457517],
    [0.2, 0.253935, 0.265254, 0.529983],
    [0.3, 0.206756, 0.371758, 0.553117],
    [0.4, 0.163625, 0.471133, 0.558148],
    [0.5, 0.127568, 0.566949, 0.550556],
    [0.6, 0.134692, 0.658636, 0.517649],
    [0.7, 0.266941, 0.748751, 0.440573],
    [0.8, 0.477504, 0.821444, 0.318195],
    [0.9, 0.741388, 0.873449, 0.150586],
    [1.0, 0.993248, 0.906157, 0.143936]
])

def _single_val_map(gray_val):
    x = gray_val / 255.0
    # 定位值所在的锚点区间
    idx = np.clip(np.searchsorted(viridis_anchors[:,0], x) - 1, 0, len(viridis_anchors)-2)
    x0, r0, g0, b0 = viridis_anchors[idx]
    x1, r1, g1, b1 = viridis_anchors[idx+1]
    # 线性插值计算通道值
    w = (x - x0) / (x1 - x0)
    return (r0 + w*(r1-r0), g0 + w*(g1-g0), b0 + w*(b1-b0), 1.0)

def manual_viridis(gray_arr):
    vec_map = np.vectorize(_single_val_map, otypes=[np.float32]*4)
    r, g, b, a = vec_map(gray_arr)
    return np.stack([r,g,b,a], axis=-1)

注：上述手动实现因为锚点是采样简化版本，和官方结果误差在1%以内；如果替换为256个全量控制点，结果可以和官方完全一致。如果需要保存为常规8位图像，将输出数组乘255后转换为np.uint8类型即可。

内容的提问来源于stack exchange，提问作者cj91