Skip to content

OpenCV 中使用 AWB

来自本人笔记: https://github.com/masterAllen/LearnOpenCV/blob/main/docs/7.2.md

OpenCV 的 xphoto 模块有三个类负责白平衡：

1. SimpleAWB：每个通道单独拉伸到固定的范围，有参数 p 表示前后 p% 的像素大小值不被拉伸时考虑
2. GrayWorldAWB：认为红绿蓝在一张图上应该差不多，根据这个假设计算出红和蓝应该分别乘以多少，有一个参数 s 表示像素的 saturation 超过 s 时不被考虑。像素的 saturation 定义为：(max(rgb)-min(rgb))/max(rgb)。
3. LearningBasedAWB: 从名字也可以看出这是通过数据驱动的方法，具体原理不谈，在该论文中。

演示代码在 test_awb.ipynb 中。

model_0 = cv2.xphoto.createSimpleWB()
# ignores the top and bottom p% of pixel values
model_0.setP(5) 

model_1 = cv2.xphoto.createGrayworldWB()
# 对于像素 i，s[i] < 0.95 才考虑，其中 s[i]=(rgb[i].max - rgb[i].min)/(rgb[i].max)
model_1.setSaturationThreshold(0.8) 

# 可传入 String 参数，表示 .yml 文件路径；不传就用默认参数
model_2 = cv2.xphoto.createLearningBasedWB()

model_0.balanceWhite(src)
model_1.balanceWhite(src)
model_2.balanceWhite(src)

LearningBasedAWB 的使用

参考来自 OpenCV 的官方文档：https://docs.opencv.org/4.x/dc/dcb/tutorial_xphoto_training_white_balance.html

我的实验数据在 ../code/image/awb 中。

1. 数据准备

数据是 Gehler-Shi dataset(http://www.cs.sfu.ca/~colour/data/shi_gehler) ，它是由两台单反相机拍摄的共 568 张图片。每张图片拍的时候都有一张色卡，这样就可以根据色卡算出 Ground Truth 了，结果在该网站的 real_illum_568..mat 文件中。

他的文件结构有点奇怪，解压之后要到很深的地方才能找到 PNG。有四个压缩包，从名字也能看出，第一个是第一个相机，其他三个是第二个相机。把这四个压缩包的文件都放在同一个目录。

2. 训练数据

使用官方的脚本文件

python learn_color_balance.py -i <path to the folder with training images> -g <path to real_illum_568..mat> -r 0,378 --num_trees 30 --max_tree_depth 6 --num_augmented 0

-r 表示训练用的图片范围，上面的目录用了 0-378 张；num_tress 和 max_tree_depth 决策树的参数，方法本质上用的是决策树。

num_augmented 表示对每张图片做多少次数据增强。一般增强都是用扭曲拉伸，但 AWB 应用场景比较特殊：他是对颜色通道各自乘以一个随机数，下面就是其中涉及的代码：

for iter in range(int(args.num_augmented)):
    # 每一次 RGB 会各自乘以随机数
    R_coef = random.uniform(0.2, 5.0); G_coef = random.uniform(0.2, 5.0); B_coef = random.uniform(0.2, 5.0)
    im_8bit = im
    im_8bit[:,:,0] *= B_coef
    im_8bit[:,:,1] *= G_coef
    im_8bit[:,:,2] *= R_coef
    im_8bit = convert_to_8bit(im)
    cur_img_features = inst.extractSimpleFeatures(im_8bit, None)
    features.append(cur_img_features.tolist())

    # gt_illumniants 就是 GroundTruth，表示颜色通道的光强大小；也需要各自乘以刚才的对应随机数
    illum = base_gt_illuminants[i]
    illum[0] *= R_coef
    illum[1] *= G_coef
    illum[2] *= B_coef
    gt_illuminants.append(illum.tolist())

一些代码上其他的点：你会在上面脚本文件中看到如下代码。这一段意思是减去黑电平，但是用的两款相机，所以黑电平不一样，分别是 0 和 129。每次图片处理需要先减去黑电平。

def load_ground_truth(gt_path):
    ...
    #Gehler-Shi dataset format
    base_gt_illuminants = gt["real_rgb"]
    # 这里就是黑电平，87 是因为第一款相机拍了 87 张，剩下的是第二款相机，他们各自黑电平分别是 0 和 129
    black_levels = 87 * [0] + (len(base_gt_illuminants) - 87) * [129]
    ...

3. 测试数据

同样使用官方的脚本

python color_balance_benchmark.py -a grayworld,learning_based:color_balance_model.yml -m <full path to folder containing the model> -i <path to the folder with training images> -g <path to real_illum_568..mat> -r 379,567 -d "img"

learning_based:color_balance_model.yml 表示训练好的模型，上面训练数据产生的模型就是这个名字，如果改了需要将这个命令也改一下。-d 表示生成的图片放在哪个目录。

实测下来，这个脚本文件还是有点问题，过曝处理的太差了，很容易偏色，我进行了如下修改，基本就没问题了：

# 266 行 main 方法中进行处理的图片先不转为 8 bit，因为会损失信息
else:
    # im = stretch_to_8bit(im)
    im = im

# 因为上面的图片没有转为 8 bit，所以 evaluate 方法需要简单修改一下
def evaluate(im, algo, gt_illuminant, i, range_thresh, bin_num, dst_folder, model_folder):
    new_im = None
    if algo=="grayworld":
        # 如果原来是 16 bit，那么就按照 16 bit 来做
        # new_im = stretch_to_8bit(im, 0)
        new_im = im.clip(0, None).astype(np.uint16)
        # ...
        new_im = inst.balanceWhite(new_im)
    elif algo=="nothing":
        new_im = im
    elif algo.split(":")[0]=="learning_based":
        new_im = im.clip(0, None).astype(np.uint16)
        # ...
        new_im = inst.balanceWhite(new_im)
    elif algo=="GT":
        new_im = im.clip(0, None).astype(float)
        g1 = gt_illuminant[1] / gt_illuminant[0]
        g3 = gt_illuminant[1] / gt_illuminant[2]

        new_im[..., 0] *= g3
        new_im[..., 2] *= g1

4. 测试其他图片

但实际把训练的模型用在其他的图片，感觉很一般... 只对这个数据集有效...

Comments