线性回归实现颜色校正

在写方程和代码之前还是先上定妆照吧，不然看了一大堆方程和代码，结果垃圾这不是耍流氓嘛~~
使用的标准色卡：
在这里插入图片描述
手机拍摄的色卡和校正后的图像：

可以看到校正后的图像颜色明显鲜艳了许多，不过白色的偏差略大了点，不过够用了。so…… 看看方程还是值得的哈哈
======================= 我是人见人烦的线性方程组分割线 =================
设第i个标准色块的RGB值为 ${R}_i ， {G}_i， {B}_i$ , 则回归方程如下：
$\begin{matrix} & {R}_i = {a}_{11}{v}_{1i} + {a}_{12}{v}_{2i} +…… + {a}_{1n}{v}_{1i}\\ & {G}_i = {a}_{21}{v}_{1i} + {a}_{22}{v}_{2i} +…… + {a}_{2n}{v}_{1i}\\ & {B}_i = {a}_{31}{v}_{1i} + {a}_{32}{v}_{2i} +…… + {a}_{3n}{v}_{1i} \end{matrix}$
有
A = $\begin{bmatrix} {a}_{11} & {a}_{12} & {a}_{13}\\ {a}_{21} & {a}_{22} & {a}_{23}\\ .....&.....&.....\\ {a}_{n1} & {a}_{n2} & {a}_{n3} \end{bmatrix}$

若 ${v}_{1i} ，{v}_{2i} ， {v}_{3i}$ 直接取原始图像的R,G,B值，回归效果不是很好，一般取R,G,B的多项式组合，即：
[1, R, G, B, RG, RB, BG, RR, BB, GG]，所以有
V = $\begin{bmatrix} {v}_{11} & {v}_{12} & ....&{v}_{1n}\\ {v}_{21} & {v}_{22} & ....&{v}_{2n}\\ .....&.....&....&.....&\\ {v}_{m1} & {v}_{m2} & ....&{v}_{mn} \end{bmatrix}$
其中，n = 10，色卡上共24个色块，所以m=24
标准色卡24个色块，所以标准色卡的RGB就是一个(3,24)的矩阵
Y = $\begin{bmatrix} {R}_{1} & {R}_{2} & ....&{R}_{24}\\ {G}_{1} & {G}_{2} & ....&{G}_{24}\\ {B}_{1} & {v}_{2} & ....&{B}_{24} \end{bmatrix}$
即： $Y = {A}^T \cdot V$
Y是标准色卡上的颜色矩阵，V是用相机拍的色卡图片颜色矩阵，A就是我们要求的系数矩阵！
根据最小二乘法优化得：
$A = {(V\cdot {V}^T)}^{-1}\cdot(V\cdot {Y}^T)$
得到A后，在校正图片时，直接使用 ${A}^T\cdot V$ 即可。搞定，好像方程也没辣么烦
================ 代码实现的分割线 ======================
理论其实很简单，不过实现起来还是有点困难的。主要是，色卡上色块的定位，需要将拍摄的图片上的色块和标准色卡上的色块顺序一致才行，不然求出来的A肯定有问题；opencv使用的是BGR顺序，这个需要特别小心；求解A的时候，只需要每个色块的一个颜色，而拍摄出来的色块肯定不止一个值，我是用色块的中间一个(5,5)区域的均值来计算的；还有最后一个问题， ${A}^T\cdot V$ 算出来的类型是float值范围远超[0,255]需要手动截断。
图像色块检测代码color_detect.py，程序中已经都注释过了，就不再解释

# -*- coding:utf-8 -*-
import cv2
import numpy as np


def _img_split_with_shadow(gray_img, threshold_value=180):
    """
    :param binary_img: 读入的灰度图
    :param img_show:
    :return: 水平和垂直线的坐标集合
    """
    h = gray_img.shape[0]
    w = gray_img.shape[1]

    # 按行求和
    sum_x = np.sum(gray_img, axis=1)
    # 按列求和
    sum_y = np.sum(gray_img, axis=0)

    h_line_index = np.argwhere(sum_x == 0)
    v_line_index = np.argwhere(sum_y == 0)

    h_line_index = np.reshape(h_line_index, (h_line_index.shape[0],))
    v_line_index = np.reshape(v_line_index, (v_line_index.shape[0],))

    h_line = []
    v_line = []

    for i in range(len(h_line_index) - 1):
        if h_line_index[i + 1] - h_line_index[i] > 2:
            h_line.append((0, h_line_index[i + 1], w - 1, h_line_index[i + 1]))
            h_line.append((0, h_line_index[i], w - 1, h_line_index[i]))

    for i in range(len(v_line_index) - 1):
        if v_line_index[i + 1] - v_line_index[i] > 2:
            v_line.append((v_line_index[i + 1], 0, v_line_index[i + 1], h - 1))
            v_line.append((v_line_index[i], 0, v_line_index[i], h - 1))

    return h_line, v_line


def _combine_rect(h_lines, v_lines):
    """
    :param h_lines: 平行直线集合
    :param v_lines: 垂直直线集合
    :return: 返回由 h_lines 和 v_lines 组成的矩形集合
    """
    rects = []

    x_axis = sorted(set([item[0] for item in v_lines]))
    y_axis = sorted(set([item[1] for item in h_lines]))

    point_list = []
    for y in y_axis:
        point = []
        for x in x_axis:
            point.append((y, x))
        point_list.append(point)

    for y_index in range(len(y_axis) - 1):
        for x_index in range(len(x_axis) - 1):
            area = abs((y_axis[y_index + 1] - y_axis[y_index]) * (x_axis[x_index + 1] - x_axis[x_index]))
            rects.append([(y_axis[y_index], x_axis[x_index],
                          y_axis[y_index + 1], x_axis[x_index + 1]), area])
    # 按面积降序排序
    rects.sort(key = lambda  ele: ele[1], reverse=True)
    areas = [ele[1] for ele in rects]

    # 找到相邻差值最大的序号
    max = -1
    index = 0
    for i in range(len(areas) - 1):
        dif = areas[i] - areas[i + 1]
        if max < dif:
            max = dif
            index = i + 1

    # rects 按坐标升序排序，使得颜色顺序和标准色卡一致
    rect_list = [ele[0] for ele in rects[0:index]]
    rect_list.sort(key = lambda  ele: ele[1])
    rect_list.sort(key = lambda ele: ele[0])

    # for i in range(len(rect_list) - 1):
    #     for j in range(0, len(rect_list) - 1 - i):
    #         if rect_list[j + 1][1] < rect_list[j][1] :
    #             rect_list[j], rect_list[j + 1] = rect_list[j + 1], rect_list[j]
    #
    # for i in range(len(rect_list) - 1):
    #     for j in range(0, len(rect_list) - 1 - i):
    #         if rect_list[j + 1][0] < rect_list[j][0]:
    #             rect_list[j], rect_list[j + 1] = rect_list[j + 1], rect_list[j]

    return rect_list


def img_split(img, img_show=False):
    """
    分割待测试的色卡图像，返回分割后的矩形图像列表和回归方程所需要的输入图像 shape:(4,6,3),像素格式：(b,g,r)
    :param img_file: 待测试色卡图像
    :param img_show: 是否显示
    :return: 分割后的子图像rect列表
    """
    # 四周各填充10个像素
    padding = 10
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 5, 7)
    binary = cv2.blur(binary, (5, 5))
    binary = cv2.bitwise_not(binary)
    binary = cv2.copyMakeBorder(binary, padding, padding, padding, padding, cv2.BORDER_CONSTANT, value=(0, 0, 0))
    # cv2.imshow('cece', binary)
    # cv2.waitKey()
    h = img.shape[0]
    w = img.shape[1]
    rate = h // w if h > w else w // h

    h_line_shadow, v_line_shadow = _img_split_with_shadow(binary)
    h_line = h_line_shadow
    v_line = v_line_shadow
    rects = _combine_rect(h_line, v_line)
    split_imgs = []

    # padding过，所以定位的时候需要减去padding的值
    img = cv2.copyMakeBorder(img, padding, padding, padding, padding, cv2.BORDER_CONSTANT, value=(0, 0, 0))
    color_img = np.zeros((4,6,3),dtype=np.uint8)
    for index, rect in enumerate(rects):
        rect_img = img[rect[0]:rect[2], rect[1]:rect[3]]
        color_img[index//6][index%6] = get_center_color(rect_img)
        # print(index, color_img[index//6][index%6])
        split_imgs.append(rect_img)

    if img_show:
        p = 0
        for rect in rects:
            cv2.rectangle(img, (rect[1], rect[0]), (rect[3], rect[2]), (0, 255, 0), 2)
            # 给识别对象写上标号
            font = cv2.FONT_HERSHEY_SIMPLEX
            cv2.putText(img, str(p), (rect[1] - 10, rect[0] + 10), font, 1, (0, 0, 255), 2)  # 加减10是调整字符位置
            p += 1

        img = cv2.resize(img, (int(h * 0.7), int(h * 0.7 / rate)))
        cv2.imshow('cece', img)
        cv2.waitKey()

    return split_imgs, color_img

def get_center_color(img):
    """
    计算给定图像中间（5，5）像素的均值
    :param img:
    :return:
    """
    w = img.shape[0]
    w = w//2
    h = img.shape[1]
    h = h//2
    data = img[h - 2:h + 2, w - 2:w + 2]
    b,g,r = cv2.split(data)
    return (int(np.mean(b)), int(np.mean(g)), int(np.mean(r)))

回归代码，也是这里的主函数

from color_detect import *

std_color_file = r'E:\code\collor_recorrect\color_value.csv'

def get_A_matrix(x, y):
    """

    :param x: 输入数据,shape:(10, n)
    :param y: 样本的标准数据, shape:(3,n)
    :return: 返回训练好的系数矩阵A， shape: (3 , 10)
    """
    temp1 = np.dot(x,x.T)
    temp2 = np.dot(x, y.T)
    temp1 = np.linalg.inv(temp1)
    A = np.dot(temp1, temp2)
    return A.T

def get_polynomial(R, G, B):
    """

    :param rgb: 像素点的RGB值,格式(r,g,b)
    :return: 返回构造的多项式，（1，R, G, B, RG, RB, BG, R*R, B*B, G*G）
    """
    R = int(R)
    G = int(G)
    B = int(B)
    return [1, R, G, B, R*G, R*B, B*G, R*R, B*B, G*G]

def create_inputData(image_data):
    """

    :param image_data: 待校正的原始图片
    :return: 返回线性回归需要的输入矩阵, shape:(10, image_data.shape[0] * image_data.shape[1])
    """
    data = []
    for raw_data in image_data:
        for bgr in raw_data:
            data.append(get_polynomial(bgr[2], bgr[1], bgr[0]))

    data = np.array(data)

    return data.T

def get_stdColor_value():
    """
    构造标准色卡的R,G,B矩阵，shape: (3 , 24)
    :return: 返回标准色卡的R,G,B值，分别用字典和矩阵存储
    """
    color_dict = {}
    std_matrix = []
    color_value_list = np.loadtxt(std_color_file, dtype=np.str, delimiter=',')

    for element in color_value_list:
        color_dict[element[1]] = (int(element[2]), int(element[3]), int(element[4]))
        std_matrix.append([int(element[2]), int(element[3]), int(element[4])])

    std_matrix = np.array(std_matrix)
    return color_dict, std_matrix.T

def recorrect_color(raw_img, A):
    """
    用系数矩阵A对图像进行颜色校正
    :param raw_img: 原始图像
    :param A: 系数矩阵
    :return: 返回校正后的图像
    """
    w = raw_img.shape[0]
    h = raw_img.shape[1]
    input_data = create_inputData(raw_img)
    corrected_data = np.dot(A, input_data)
    data = []
    for element in corrected_data:
        vec = []
        for value in element:
            if 0.0 <= value <= 255.0:
                vec.append(int(value))
            elif 0.0 > value:
                vec.append(0)
            elif 255.0 < value:
                vec.append(255)
        data.append(vec)

    data = np.array(data)
    data = data.transpose((1, 0))
    new_img = data.reshape((w,h,3))
    cv2.imwrite(r'E:\code\collor_recorrect\correct_test.jpg', new_img[...,[2,1,0]])
    return new_img


if __name__ == '__main__':
    # 载入标准色卡数据
    color_dict, std_matrix = get_stdColor_value()
    # 载入测试色卡图像，生成回归输入数据
    img = cv2.imread(r'E:\code\collor_recorrect\data\test.jpg', 1)
    imgs, color_img = img_split(img)
    input_data = create_inputData(color_img)
    # 计算回归方程的系数矩阵
    A = get_A_matrix(input_data, std_matrix)
    # 颜色校正
    recorrect_color(img, A)

标准色卡的csv文件：

1,Dark Skin,94,28,13
2,Light Skin,241,149,108
3,Bolu Sky,97,119,171
4,Foliage,90,103,39
5,Blue Flower,164,131,196
6,Bluish Green,140,253,153
7,Orange,255,116,21
8,Purplish Blue,7,47,122
9,Moderate Red,222,29,42
10,Purple,69,0,68
11,Yellow Green,187,255,19
12,Orange Yellow,255,142,0
13,Blue,0,0,142
14,Green,64,173,38
15,Red,203,0,0
16,Yellow,255,217,0
17,Magenta,207,3,124
18,Cyan,0,148,189
19,White (.05)*,255,255,255
20,Neutral 8 (.23) *,249,249,249
21,Neutral6.5 (.44) *,180,180,180
22,Neutral 5 (.70) *,117,117,117
23,Neutral3.5 (1.05) *,53,53,53
24,Black (1.5) *,0,0,0

两天的成绩，貌似结果还行。哈哈~

线性回归实现颜色校正

猜你喜欢