4. AWB进阶

任务目标

本案例实现一个基于灰度世界与完美反射融合的色温估计器，对输入图像统计R/G/B通道均值，估计场景色温，并据此计算R/B通道增益完成白平衡校正。最终输出校正后图像及估计色温值，用于验证算法在典型光源下的表现。

核心实现

import numpy as np
import cv2

def estimate_color_temperature(img_bgr):
    """基于灰度世界+完美反射融合估计色温(K)"""
    # 分离通道
    b, g, r = cv2.split(img_bgr.astype(np.float32))
    
    # 灰度世界假设：各通道均值应趋于相等
    mean_r, mean_g, mean_b = np.mean(r), np.mean(g), np.mean(b)
    gray_world_gain_r = mean_g / (mean_r + 1e-6)
    gray_world_gain_b = mean_g / (mean_b + 1e-6)
    
    # 完美反射假设：最亮像素应接近白色(取前1%高亮像素)
    h, w = img_bgr.shape[:2]
    top_k = max(1, int(h * w * 0.01))
    # 计算亮度并取最亮像素索引
    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY).flatten()
    idx = np.argpartition(gray, -top_k)[-top_k:]
    bright_r = r.flatten()[idx].mean()
    bright_g = g.flatten()[idx].mean()
    bright_b = b.flatten()[idx].mean()
    perfect_ref_gain_r = bright_g / (bright_r + 1e-6)
    perfect_ref_gain_b = bright_g / (bright_b + 1e-6)
    
    # 融合两种假设(各占50%)
    gain_r = 0.5 * gray_world_gain_r + 0.5 * perfect_ref_gain_r
    gain_b = 0.5 * gray_world_gain_b + 0.5 * perfect_ref_gain_b
    
    # 色温估计：基于R/B增益比映射到色温范围
    rb_ratio = (gain_r / (gain_b + 1e-6))
    # 经验映射：ratio=1.0对应5500K(日光)，ratio增大对应更低色温(暖)
    # 使用对数映射使结果更平滑
    temp_k = 5500.0 - 1800.0 * np.log(rb_ratio + 1e-6)
    temp_k = np.clip(temp_k, 2500.0, 10000.0)
    
    return temp_k, gain_r, gain_b

def awb_correct(img_bgr):
    """执行白平衡校正，返回校正图与估计色温"""
    temp_k, gain_r, gain_b = estimate_color_temperature(img_bgr)
    
    # 应用增益(限制增益范围避免过曝)
    gain_r = np.clip(gain_r, 0.5, 2.5)
    gain_b = np.clip(gain_b, 0.5, 2.5)
    
    b, g, r = cv2.split(img_bgr.astype(np.float32))
    r_corr = np.clip(r * gain_r, 0, 255)
    b_corr = np.clip(b * gain_b, 0, 255)
    g_corr = g  # G通道保持不变
    
    result = cv2.merge([b_corr, g_corr, r_corr]).astype(np.uint8)
    return result, temp_k

# 最小调用示例
if __name__ == "__main__":
    # 生成模拟测试图：低色温(暖光)下拍摄的灰卡
    test_img = np.full((200, 200, 3), (200, 150, 100), dtype=np.uint8)  # BGR暖色调
    corrected, temp = awb_correct(test_img)
    print(f"估计色温: {temp:.0f}K")
    print(f"校正后像素(BGR): {corrected[100, 100].tolist()}")
    # 预期输出: 估计色温约3000-4000K(暖光), 校正后像素接近灰色(各通道值接近)

运行与验证

依赖环境：Python 3.8+，NumPy，OpenCV（pip install numpy opencv-python）。

运行方式：将上述代码保存为 awb_estimation.py，直接执行 python awb_estimation.py。

验证方法：

1. 观察终端输出的估计色温值：对于暖色输入（B>G>R），色温应显著低于5500K（约2500-4000K区间）；
2. 检查校正后像素：原暖色像素（如B=200,G=150,R=100）经校正后，B/R通道值应趋近G通道值（约150±20），证明白平衡生效；
3. 使用真实图像测试：拍摄一张在钨丝灯（约2800K）下的白纸照片，传入函数后校正结果应接近中性灰，且估计色温在2500-3500K范围内。

通过标准：校正后图像中原本偏色的白色区域RGB三通道最大差值小于30，且估计色温与真实光源色温偏差不超过±500K（在标准光源下测试时）。

关键说明

本实现采用灰度世界与完美反射加权融合，避免单一假设在纯色场景下失效。色温映射公式基于R/B增益比的经验对数关系，覆盖2500K-10000K典型范围。增益限制在0.5-2.5倍防止极端校正导致色彩失真。对于包含大面积单一颜色的场景，估计精度会下降，这是统计方法的固有局限。