02_ccm_useLstsq

复制本地路径 | 在线编辑
原始文件为 .py 代码，本文是转换后的 Markdown 文件。
'''
进行颜色矫正，用最小二乘法计算 CCM
'''
import os
os.chdir(os.path.dirname(os.path.abspath(__file__)))

import numpy as np
from func.cp_hw2 import read_colorchecker_gm, readHDR, writeHDR

# 该函数用于生成色卡块图，用于检测正确性；同时也顺便生成色卡每个块的平均值
def make_patches(coors, img, onerow=6):
    rows, cols = 24//onerow, onerow

    # patches: 色卡的图，patches_rgb：色卡的平均图(24*3)
    patches = np.zeros((rows*80, cols*80, 3))
    patches_rgb = np.zeros((24, 3))
    for i, coor in enumerate(coors):
        nowpatch = img[coor[1]:coor[3], coor[0]:coor[2]]

        x = (i // onerow) * 80
        y = (i % onerow) * 80
        patches[x:x+80, y:y+80] = nowpatch[0:80, 0:80]
        patches_rgb[i] = np.mean(nowpatch[0:80, 0:80].reshape(-1, 3), axis=0)
    return patches, patches_rgb

coors = []
with open('../data/24color_coor.txt', 'r') as f:
    for i, line in enumerate(f):
        coors.append([int(x) for x in line.split()])

# 保存的色卡坐标，索引有点不对，现在重排一下
coors = np.array(coors).reshape(6, 4, 4).transpose(1, 0, 2).reshape(24, 4)

# 加载基准 RGB
ok_rgb, ok_Lab = read_colorchecker_gm()

# 开始处理各个图片
os.makedirs('../result/stage-2', exist_ok=True)
for filename in os.listdir('../result/stage-1'):
    print(filename)

    hdrimg = readHDR(f'../result/stage-1/{filename}')

    # 先 AWB：提取色卡，只用白色块进行校准，最后一行第一个
    now_patches, patches_rgb = make_patches(coors, hdrimg)
    wr, wg, wb = patches_rgb[3*6+0]

    rgain, bgain = wg / wr, wg / wb
    hdrimg[..., 0] *= rgain
    hdrimg[..., 2] *= bgain


    # 根据coors记录的色卡坐标，生成色卡图，顺便把色卡的各个平均值计算一下
    now_patches, patches_rgb = make_patches(coors, hdrimg)

    # 计算 CCM 矩阵
    # Au = ok_rgb, A->26*4, u->4*3， ok_rgb->26*3
    A = np.hstack((patches_rgb, np.ones((24, 1))))
    u, *_ = np.linalg.lstsq(A, ok_rgb, rcond=None)

    # Make CCM
    rows, cols = hdrimg.shape[0:2]
    hdrimg = np.hstack((hdrimg.reshape(-1, 3), np.ones((rows*cols, 1))))
    hdrimg = np.dot(hdrimg.astype(np.longdouble), u.astype(np.longdouble))

    # ! 很重要：否则会有异色！！！
    hdrimg[hdrimg < 0] = 0
    hdrimg = hdrimg.astype(np.float32).reshape((rows, cols, 3))

    # WriteHDR, Good Job.
    writeHDR(f'../result/stage-2/{filename}', hdrimg)
02_ccm_useLstsq

Comments
