如何快速实现基于最近邻分类器的图像/视频像素级分类处理？

问题背景 #

当前通过最近邻分类器实现图像逐像素分类：手动指定RGB质心与阈值，计算每个像素与质心的欧氏距离，超过阈值则设为黑色，否则保留原值。原代码因Python嵌套循环导致单帧处理耗时2-3秒，无法满足实时视频10fps的需求。

核心优化思路 #

Python嵌套循环是性能瓶颈，改用Numpy向量化操作（底层为C实现）替代逐像素循环，同时通过数学简化减少计算量，大幅提升处理速度。

优化后的代码 #

import cv2
import numpy as np
import matplotlib.pyplot as plt

# 读取图像，注意OpenCV默认BGR格式
frame = cv2.imread('coloredChips.png')
# 复制图像避免修改原数据
mask = frame.copy()

# 配置参数
threshold = 100
# 质心的RGB值（对应OpenCV的BGR通道顺序：cen_b对应通道0，cen_g对应1，cen_r对应2）
cen_b, cen_g, cen_r = 10, 10, 111

# 向量化计算每个像素与质心的欧氏距离平方（避免开根号，减少计算量）
diff = frame - np.array([cen_b, cen_g, cen_r], dtype=np.uint8)
dist_sq = np.sum(diff ** 2, axis=2)
# 阈值平方，和距离平方比较结果等价
threshold_sq = threshold ** 2

# 将距离超过阈值的像素设为黑色
mask[dist_sq > threshold_sq] = [0, 0, 0]

plt.imshow(cv2.cvtColor(mask, cv2.COLOR_BGR2RGB))  # 转换为RGB格式显示
plt.show()

优化细节说明 #

• 替换嵌套循环：用Numpy的数组广播机制一次性计算所有像素的差值，避免Python循环的性能开销，速度提升可达数十倍。
• 简化距离计算：比较欧氏距离与阈值等价于比较距离平方与阈值平方，省去np.sqrt或np.linalg.norm的开根号运算，进一步减少计算量。
• 避免原图像修改：原代码中mask=frame是引用赋值，会直接修改原始图像，改用frame.copy()创建独立副本。
• 通道顺序对齐：OpenCV读取图像为BGR格式，将质心参数按BGR顺序定义，避免重复的通道索引转换。

实时视频处理适配 #

若要适配实时视频，只需将单帧处理逻辑封装到循环中，例如：

cap = cv2.VideoCapture(0)  # 0为默认摄像头

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    
    # 复用上述优化后的处理逻辑
    diff = frame - np.array([cen_b, cen_g, cen_r], dtype=np.uint8)
    dist_sq = np.sum(diff ** 2, axis=2)
    mask = frame.copy()
    mask[dist_sq > threshold_sq] = [0, 0, 0]
    
    cv2.imshow('Processed Frame', mask)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

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