什么是图像白平衡？

肯定有人会问什么是白平衡，这里给大家简单说一下。白平衡改变的是画面的色温冷暖变化，要么让画面偏蓝（冷），要么让画面偏黄（暖）。

原图

白平衡处理后（画面偏冷）

涉及难点

策略模式：在有多种算法相似的情况下，使用“if...else...”会带来的难以维护的问题。为了避免过多使用多重条件的判断，而且为了扩展性良好，就得选择策略模式。

递归访问文件：考虑到用户给出的待处理的图片可能有含有多级目录，递归访问能确保处理好所有的图片。

numpy的运用：也可以使用数组代替，然后用for循环逐个像素值进行计算。不过，考虑的运行效率问题，最好选择numpy或pandas。

python魔术方法：可以这么说，就是函数的声明和调用所使用的名称不同，而且调用的方式非常方便。

项目设计

首先申明，本代码是根据网上开源的代码进行修改的。加入了python的设计理念和批量处理的思路。主要分为七步，大概有200行python代码，大家可以研读一下。

• 第一步：导入需要的包和静态变量的设置。

import cv2
import numpy as np
import os


class CONF:
    input_path = 'input'    # 待处理的图片（路径不能有中文）
    output_path = 'output'  # 处理的结果（路径不能有中文）
    white_balance_mode = 5  # 白平衡的模式（一共有五种模式，1~5）
    is_log = True           # 是否打印日志

• 第二步：类的初始化

def __init__(self, input_path, output_path, white_balance_mode, is_log):


    # 策略的选择
    self.white_balance_mode = white_balance_mode if white_balance_mode in [1,2,3,4,5] else 1
    self.name_dict = {1:'_mode_1', 2:'_mode_2', 3:'_mode_3', 4:'_mode_4', 5:'_mode_5'}
    self.input_path = input_path
    self.output_path = output_path + self.name_dict[self.white_balance_mode]


    # 策略模式
    self.white_balance_process = {1: self.__balance_mode_1, # 均值
                                  2: self.__balance_mode_2, # 完美反射
                                  3: self.__balance_mode_3, # 灰度世界
                                  4: self.__balance_mode_4, # 颜色校正
                                  5: self.__balance_mode_5} # 动态阈值
    self.is_log = is_log


"""初始化"""
@classmethod
def initialize(cls, config):
    input_path = config.input_path
    output_path = config.output_path
    white_balance_mode = config.white_balance_mode
    is_log = config.is_log
    return cls(input_path, output_path, white_balance_mode, is_lo

• 第三步：五种策略算法的编写。

策略一：均值

"""策略一： 均值"""
def __balance_mode_1(self, img, b, g, r, h, w, c):
    # 简单的求均值白平衡法
    b_avg, g_avg, r_avg = cv2.mean(b)[0], cv2.mean(g)[0], cv2.mean(r)[0]
    # 求各个通道所占增益
    k = (b_avg + g_avg + r_avg) / 3
    kb, kg, kr = k / b_avg, k / g_avg, k / r_avg
    b = cv2.addWeighted(src1=b, alpha=kb, src2=0, beta=0, gamma=0)
    g = cv2.addWeighted(src1=g, alpha=kg, src2=0, beta=0, gamma=0)
    r = cv2.addWeighted(src1=r, alpha=kr, src2=0, beta=0, gamma=0)
    output_img = cv2.merge([b, g, r])
    return output_img

策略二：完美反射

"""策略二： 完美反射"""
def __balance_mode_2(self, img, b, g, r, h, w, c):
    # 完美反射白平衡 ---- 依赖ratio值选取而且对亮度最大区域不是白色的图像效果不佳。
    output_img = img.copy()
    sum_ = np.double() + b + g + r
    hists, bins = np.histogram(sum_.flatten(), 766, [0, 766])
    Y = 765
    num, key = 0, 0
    ratio = 0.01
    while Y >= 0:
        num += hists[Y]
        if num > h * w * ratio / 100:
            key = Y
            break
        Y = Y - 1


    sumkey = np.where(sum_ >= key)
    sum_b, sum_g, sum_r = np.sum(b[sumkey]), np.sum(g[sumkey]), np.sum(r[sumkey])
    times = len(sumkey[0])
    avg_b, avg_g, avg_r = sum_b / times, sum_g / times, sum_r / times


    maxvalue = float(np.max(output_img))
    output_img[:, :, 0] = output_img[:, :, 0] * maxvalue / int(avg_b)
    output_img[:, :, 1] = output_img[:, :, 1] * maxvalue / int(avg_g)
    output_img[:, :, 2] = output_img[:, :, 2] * maxvalue / int(avg_r)
    return output_im

策略三：灰度世界

"""策略三： 灰度世界"""
def __balance_mode_3(self, img, b, g, r, h, w, c):
    # 灰度世界假设
    b_avg, g_avg, r_avg = cv2.mean(b)[0], cv2.mean(g)[0], cv2.mean(r)[0]
    # 需要调整的RGB分量的增益
    k = (b_avg + g_avg + r_avg) / 3
    kb, kg, kr = k / b_avg, k / g_avg, k / r_avg
    ba, ga, ra = b * kb, g * kg, r * kr
    output_img = cv2.merge([ba, ga, ra])
    return output_img

策略四：颜色校正

"""策略四： 颜色校正"""
def __balance_mode_4(self, img, b, g, r, h, w, c):
    # 基于图像分析的偏色检测及颜色校正
    I_b_2, I_r_2 = np.double(b) ** 2, np.double(r) ** 2
    sum_I_b_2, sum_I_r_2 = np.sum(I_b_2), np.sum(I_r_2)
    sum_I_b, sum_I_g, sum_I_r = np.sum(b), np.sum(g), np.sum(r)
    max_I_b, max_I_g, max_I_r = np.max(b), np.max(g), np.max(r)
    max_I_b_2, max_I_r_2 = np.max(I_b_2), np.max(I_r_2)
    [u_b, v_b] = np.matmul(np.linalg.inv([[sum_I_b_2, sum_I_b], [max_I_b_2, max_I_b]]), [sum_I_g, max_I_g])
    [u_r, v_r] = np.matmul(np.linalg.inv([[sum_I_r_2, sum_I_r], [max_I_r_2, max_I_r]]), [sum_I_g, max_I_g])
    b0 = np.uint8(u_b * (np.double(b) ** 2) + v_b * b)
    r0 = np.uint8(u_r * (np.double(r) ** 2) + v_r * r)
    output_img = cv2.merge([b0, g, r0])
    return output_img

策略五：动态阈值

"""策略五： 动态阈值"""
def __balance_mode_5(self, img, b, g, r, h, w, c):
    # 动态阈值算法 ---- 白点检测和白点调整
    # 只是白点检测不是与完美反射算法相同的认为最亮的点为白点，而是通过另外的规则确定
    def con_num(x):
        if x > 0:
            return 1
        if x < 0:
            return -1
        if x == 0:
            return 0


    yuv_img = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
    # YUV空间
    (y, u, v) = cv2.split(yuv_img)
    max_y = np.max(y.flatten())
    sum_u, sum_v = np.sum(u), np.sum(v)
    avl_u, avl_v = sum_u / (h * w), sum_v / (h * w)
    du, dv = np.sum(np.abs(u - avl_u)), np.sum(np.abs(v - avl_v))
    avl_du, avl_dv = du / (h * w), dv / (h * w)
    radio = 0.5  # 如果该值过大过小，色温向两极端发展


    valuekey = np.where((np.abs(u - (avl_u + avl_du * con_num(avl_u))) < radio * avl_du)
                        | (np.abs(v - (avl_v + avl_dv * con_num(avl_v))) < radio * avl_dv))
    num_y, yhistogram = np.zeros((h, w)), np.zeros(256)
    num_y[valuekey] = np.uint8(y[valuekey])
    yhistogram = np.bincount(np.uint8(num_y[valuekey].flatten()), minlength=256)
    ysum = len(valuekey[0])
    Y = 255
    num, key = 0, 0
    while Y >= 0:
        num += yhistogram[Y]
        if num > 0.1 * ysum:  # 取前10%的亮点为计算值，如果该值过大易过曝光，该值过小调整幅度小
            key = Y
            break
        Y = Y - 1


    sumkey = np.where(num_y > key)
    sum_b, sum_g, sum_r = np.sum(b[sumkey]), np.sum(g[sumkey]), np.sum(r[sumkey])
    num_rgb = len(sumkey[0])


    b0 = np.double(b) * int(max_y) / (sum_b / num_rgb)
    g0 = np.double(g) * int(max_y) / (sum_g / num_rgb)
    r0 = np.double(r) * int(max_y) / (sum_r / num_rgb)


    output_img = cv2.merge([b0, g0, r0])
    return output

• 第四步： 单张图片的处理

"""单张图片的处理"""
def img_process(self, in_img_path, out_img_path):


    in_img = cv2.imread(in_img_path)


    b, g, r = cv2.split(in_img)
    h, w, c = in_img.shape  # 均值变为三通道


    # 根据用户选择的不同策略进行调用
    output_img = self.white_balance_process[self.white_balance_mode](in_img, b, g, r, h, w, c)
    output_img = np.uint8(np.clip(output_img, 0, 255))


    # 保存图片
    cv2.imwrite(out_img_path, output_i

• 第五步：递归访问文件

""" 创建文件夹 """
def mkdir(self, path):
    path = path.strip().rstrip("\\")
    is_exists = os.path.exists(path)
    if not is_exists:
        os.makedirs(path)
        return True
    else:
        return False




""" 递归访问文件/文件夹 """
def visit_dir_files(self, org_input_dir, org_output_dir, recursion_dir):
    single_file = False
    if os.path.isdir(recursion_dir):
        dir_list = os.listdir(recursion_dir)
    else:
        dir_list = [recursion_dir]
        single_file = True
    for i in range(0, len(dir_list)):
        path = os.path.join(recursion_dir, dir_list[i])
        if os.path.isdir(path):
            self.visit_dir_files(org_input_dir, org_output_dir, path)
        else:
            abs_output_dir = org_output_dir + recursion_dir[len(org_input_dir):]
            target_path = os.path.join(abs_output_dir, dir_list[i])
            if single_file:
                target_path = os.path.join(org_output_dir, os.path.basename(dir_list[i]))
            target_dir_name = os.path.dirname(target_path)
            if not os.path.exists(target_dir_name):
                self.mkdir(target_dir_name)
            self.img_process(path, target_path

• 第六步：主线程，这里使用魔术方法声明。

# 主线程(采用python魔术方法实现)
# 允许一个类的实例像函数一样被调用：x(key)
# 假如对象实例是x, 使用x(key), 就像调用 x.__call__(a, b)一样
def __call__(self,):
    if os.path.exists(self.input_path):
        # 递归访问图片，逐张图片进行处理
        self.visit_dir_files(self.input_path, self.output_path, self.input_path)
        if self.is_log:
            print(u'完成！所有图片已保存至路径' + self.output_path)
    else:
        print(u'待处理的图片存放的位置 %s, 如果没有请新建目录 %s' % (self.input_path, self.input_path))

• 第七步：main函数进行调用

if __name__ == '__main__':


    white_balance = WhiteBalanceHelper.initialize(config=CONF)
    white_balance()

输出结果

第一张是原图，余下五张是不同白平衡算法的处理结果。

后续改进

项目采用了五种白平衡的算法，并用策略模式规范起来，目的就是为了后续方便扩展其他算法。如果有其他新的白平衡算法，可以逐步加入，增加变量 mode的取值以及策略字典的函数值即可。