如何利用Numpy数组操作避免for循环实现热图像素色标匹配

某设备可输出附带色标比例尺的热图图像，但不支持将数据导出为.csv等可直接导入Python分析的格式。
已通过PIL完成图像导入，得到以下数组：

• frame1：维度为(680, 900, 3)的热图数组，存储每个像素的3通道RGB值，对应热图的XY空间坐标
• scale1：维度为(254, 3)的色标采样数组，存储色标上254个采样点的3通道RGB值
• scaleval：通过np.linspace(maxval, minval, 254)生成的线性刻度序列，其中maxval、minval为从图像中读取的色标最大、最小刻度值
  需求为将frame1中每个像素匹配到scale1中最接近的颜色，再将scaleval中对应的刻度值存入数据框df。原有三层for循环实现运行速度极慢，需要无显式循环的高效Numpy实现。
  原有循环参考代码：

import math
import numpy as np
import pandas as pd

# 计算两个RGB值的欧氏距离
def distance(c1, c2):
    (r1,g1,b1) = c1
    (r2,g2,b2) = c2
    return math.sqrt((r1 - r2)**2 + (g1 - g2) ** 2 + (b1 - b2) **2)

# 遍历热图每一列
for j in range(frame1.shape[1]):
    # 遍历热图每一行
    for k in range(frame1.shape[0]):
        distances = []
        # 遍历所有色标点计算距离
        for i in range(len(scale1)):
            distances.append(distance(scale1[i], frame1[k,j,:]))
        
        distarr = np.asarray(distances)
        idx = distarr.argmin()
        df.loc[k,j] = scaleval[idx]
    print(f"Column {j+1} completed") 

欧氏距离的大小比较不需要做开平方运算——平方是单调递增函数，直接用距离平方找最小值，结果和开平方完全一致，还能省去开平方的计算开销。利用Numpy的广播机制可以一次性完成所有像素到所有色标点的距离计算，全程无Python层循环，速度比三层循环提升两个数量级以上。

纯Numpy实现（无额外依赖） #

import numpy as np
import pandas as pd

# 转为浮点型避免uint8类型计算溢出
frame = frame1.astype(np.float32)
scale = scale1.astype(np.float32)

# 广播计算所有像素到每个色标点的RGB距离平方，输出维度(680, 900, 254)
dist_sq = np.sum(
    (frame[:, :, np.newaxis, :] - scale[np.newaxis, np.newaxis, :, :]) ** 2,
    axis=-1
)

# 沿色标维度取最小距离对应的索引
closest_idx = dist_sq.argmin(axis=-1)

# 直接映射得到每个像素对应的刻度值，转为DataFrame
value_map = scaleval[closest_idx]
df = pd.DataFrame(value_map)

更快的KDTree实现（依赖scipy） #

如果对速度有更高要求，可以用KDTree做最近邻查询，色标样本量小的场景下和纯Numpy实现差异不大，大分辨率热图下速度优势更明显：

import numpy as np
import pandas as pd
from scipy.spatial import cKDTree

# 基于色标颜色构建KD树
kdtree = cKDTree(scale1)
# 批量查询所有像素的最近邻色标点
_, closest_idx = kdtree.query(frame1.reshape(-1, 3), k=1)
# 重构为热图原始尺寸，转为DataFrame
value_map = scaleval[closest_idx].reshape(frame1.shape[0], frame1.shape[1])
df = pd.DataFrame(value_map)

两种实现处理680*900分辨率的热图耗时均在100ms以内，不需要逐层打印列处理进度。

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