灰点检测系列2：利用灰点检测估计光源，求白平衡

这个方法并不是灰点检测方法做白平衡。

假如每个像素代表一种颜色，白平衡有两个思路，一个是找灰色，另一个是找更多颜色。找更多颜色什么意思呢，加入一张图片五颜六色，但是没有灰色，这个时候我们希望找到更多的颜色，这样平均一下可能会得到一个灰色，和灰度世界思想差不多吧，找到一个红色一个蓝色，两者想加可能接近灰色，颜色越多越接近灰色

这里就是选取距离灰色比较远的颜色的点，在亮暗像素之间拟合曲线

import os

import cv2
import numpy as np

import matplotlib.pyplot as plt


def illuminant_estimator(img, mask, prc):
    if mask.size != 0:
        img = img[mask, :]

    data = np.reshape(img, (-1, 3))
    n = len(data)

    # gray world
    l = np.mean(data, 0)
    l = l / np.linalg.norm(l)
    l = l.reshape(-1, 3)

    data_p = np.sum(data * l, axis=1)
    data_idx = np.argsort(data_p)

    # select pixel
    low = int(np.ceil(n * prc / 100))
    high = int(np.floor(n * (100 - prc) / 100))
    low_sel = data_idx[:low]
    high_sel = data_idx[high:]

    data_selected = np.vstack((data[low_sel].reshape(-1, 3), data[high_sel].reshape(-1, 3)))

    sigma = np.matmul(data_selected.T, data_selected)
    w, v = np.linalg.eig(sigma)
    s = np.argsort(w)[-1]
    ei = abs(v[:, s])
    return l.reshape(-1), ei


def cal_angle(light, ref):
    light = np.reshape(light, (-1, 3))
    ref = np.reshape(ref, (-1, 3))

    cos_angle = np.sum(light * ref, axis=1) / (
            np.sqrt(np.sum(np.power(light, 2), 1)) * np.sqrt(np.sum(np.power(ref, 2), 1)))

    angle = np.arccos(cos_angle)
    angle = angle * 180 / np.pi
    return angle


if __name__ == "__main__":
    # single image process

    filename = r'G:\ffcc-master_20201108\ffcc-master\data\shi_gehler\preprocessed\GehlerShi\000002.png'
    imbgr = cv2.imread(filename)
    im = imbgr[:, :, ::-1]
    gt_file = filename[:-4] + '.txt'
    ill_gt = np.loadtxt(gt_file)
    maskrgb = np.logical_or((im <= 254), (im > 0))
    mask = np.logical_and(maskrgb[..., 0], maskrgb[..., 1], maskrgb[..., 2])

    prc = 1
    im = im / 255
    gw_ill, spatial_domain_ill = illuminant_estimator(im, mask, prc)

    print('wb_gain : ', spatial_domain_ill[1] / spatial_domain_ill)

    # batch process

    dir = r'G:\ffcc-master_20201108\ffcc-master\data\shi_gehler\preprocessed\GehlerShi'
    gw_ills = []
    spatial_domain_ills = []

    refs = []

    filesets = os.listdir(dir)
    for file in filesets:
        if file.endswith('.png'):
            filename = os.path.join(dir, file)

            im = cv2.imread(filename)
            im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
            im = im / 255

            gt_file = filename[:-4] + '.txt'
            gt = np.loadtxt(gt_file)
            refs.append(gt)

            prc = 4
            gw_ill, spatial_domain_ill = illuminant_estimator(im, mask, prc)
            gw_ills.append(gw_ill)
            spatial_domain_ills.append(spatial_domain_ill)
            print('file and result :', gw_ill, spatial_domain_ill, gt)
    diff1 = cal_angle(gw_ills, refs).reshape(-1, 1)
    diff2 = cal_angle(spatial_domain_ills, refs).reshape(-1, 1)

    diff = np.hstack((diff1, diff2))
    print('angle : ', diff, '\n', np.mean(diff, axis=0))