DH-AISP/1/awb.py

import os
import cv2
import numpy as np
from glob import glob


def dynamic(rgb):

    rgb = rgb[:-1, :-1, :]  # 删去一行一列
    h, w, _ = rgb.shape
    col = 4
    row = 3
    h1 = h // row
    w1 = w // col

    r, g, b = cv2.split(rgb)
    r_mask = r < 0.95
    g_mask = g < 0.95
    b_mask = b < 0.95
    mask = r_mask * g_mask * b_mask
    r *= mask
    g *= mask
    b *= mask
    rgb = np.stack((r, g, b), axis=2)

    y, cr, cb = cv2.split(cv2.cvtColor(rgb, cv2.COLOR_RGB2YCrCb))
    cr -= 0.5
    cb -= 0.5

    mr, mb, dr, db = 0, 0, 0, 0
    for r in range(row):
        for c in range(col):
            cr_1 = cr[r * h1:(r + 1) * h1, c * w1:(c + 1) * w1]
            cb_1 = cb[r * h1:(r + 1) * h1, c * w1:(c + 1) * w1]
            mr_1 = np.mean(cr_1)
            mb_1 = np.mean(cb_1)
            dr_1 = np.mean(np.abs(cr_1 - mr))
            db_1 = np.mean(np.abs(cb_1 - mb))

            mr += mr_1
            mb += mb_1
            dr += dr_1
            db += db_1

    mr /= col * row
    mb /= col * row
    dr /= col * row
    db /= col * row

    cb_mask = np.abs(cb - (mb + db * np.sign(mb))) < 1.5 * db
    cr_mask = np.abs(cr - (1.5 * mr + dr * np.sign(mr))) < 1.5 * dr

    mask = cb_mask * cr_mask
    y_white = y * mask

    hist_y = np.zeros(256, dtype=np.int)
    y_white_uint8 = (y_white * 255).astype(np.int)

    for v in range(255):
        hist_y[v] = np.sum(y_white_uint8 == v)

    thr_sum = 0.05 * np.sum(mask)
    sum_v = 0
    thr = 0
    for v in range(255, -1, -1):
        sum_v = sum_v + hist_y[v]
        if sum_v > thr_sum:
            thr = v
            break

    white_mask = y_white_uint8 > thr
    cv2.imwrite(r'V:\Project\3_MEWDR\data\2nd_awb\t.png', (white_mask + 0) * 255)

    r, g, b = cv2.split(rgb)
    r_ave = np.sum(r[white_mask]) / np.sum(white_mask)
    g_ave = np.sum(g[white_mask]) / np.sum(white_mask)
    b_ave = np.sum(b[white_mask]) / np.sum(white_mask)

    return 1 / r_ave, 1 / g_ave, 1 / b_ave


def perf_ref(rgb, eps):
    h, w, _ = rgb.shape

    r, g, b = cv2.split(rgb)
    r_mask = r < 0.95
    g_mask = g < 0.95
    b_mask = b < 0.95
    mask = r_mask * g_mask * b_mask
    r *= mask
    g *= mask
    b *= mask
    rgb = np.stack((r, g, b), axis=2)
    rgb = np.clip(rgb * 255, 0, 255).astype(np.int)

    hist_rgb = np.zeros(255 * 3, dtype=np.int)
    rgb_sum = np.sum(rgb, axis=2)

    for v in range(255 * 3):
        hist_rgb[v] = np.sum(rgb_sum == v)

    thr_sum = eps * h * w
    sum_v = 0
    thr = 0
    for v in range(255 * 3 - 1, -1, -1):
        sum_v = sum_v + hist_rgb[v]
        if sum_v > thr_sum:
            thr = v
            break

    thr_mask = rgb_sum > thr
    r_ave = np.sum(r[thr_mask]) / np.sum(thr_mask)
    g_ave = np.sum(g[thr_mask]) / np.sum(thr_mask)
    b_ave = np.sum(b[thr_mask]) / np.sum(thr_mask)

    # k = (r_ave + g_ave + b_ave) / 3.
    # k = 255

    # print(k)

    # r = np.clip(r * k / r_ave, 0, 255)
    # g = np.clip(g * k / g_ave, 0, 255)
    # b = np.clip(b * k / b_ave, 0, 255)

    return 1 / r_ave, 1 / g_ave, 1 / b_ave


def awb_v(in_image, bayer, eps):
    
    assert bayer in ['GBRG', 'RGGB']
    
    if bayer == 'GBRG':
        g = in_image[0::2, 0::2]  # [0,0]
        b = in_image[0::2, 1::2]  # [0,1]
        r = in_image[1::2, 0::2]  # [1,0]
    else:
        r = in_image[0::2, 0::2]  # [0,0]
        g = in_image[0::2, 1::2]  # [0,1]
        b = in_image[1::2, 1::2]  # [1,1]
    
    rgb = cv2.merge((r, g, b)) * 1

    r_gain, g_gain, b_gain = perf_ref(rgb, eps)
    
    return r_gain / g_gain, b_gain / g_gain


def main():
    path = r'V:\Project\3_MEWDR\data\2nd_raw'
    # out_path = r'V:\Project\3_MEWDR\data\2nd_awb'

    files = glob(os.path.join(path, '*.png'))

    for f in files:
        img = cv2.imread(f, cv2.CV_16UC1)
        img = (img.astype(np.float) - 2048) / (15400 - 2048) * 4

        g = img[0::2, 0::2]  # [0,0]
        b = img[0::2, 1::2]  # [0,1]
        r = img[1::2, 0::2]  # [1,0]
        # g_ = img[1::2, 1::2]

        rgb = cv2.merge((r, g, b))

        # save_name = f.replace('.png', '_rgb.png').replace('2nd_raw', '2nd_awb')

        r_gain, g_gain, b_gain = perf_ref(rgb, eps=0.1)
        # r_gain, g_gain, b_gain = dynamic(rgb.astype(np.float32))

        r *= r_gain / g_gain
        b *= b_gain / g_gain
        print(r_gain / g_gain, b_gain / g_gain)

        out_rgb = np.clip(cv2.merge((r, g, b)) * 255, 0, 255)

        save_name = f.replace('.png', '_awb4_dyn.png').replace('2nd_raw', '2nd_awb')

        cv2.imwrite(save_name, cv2.cvtColor(out_rgb.astype(np.uint8), cv2.COLOR_RGB2BGR))

        # break


if __name__ == '__main__':
    main()