opencv/samples/dnn/auto_white_balance.py

#!/usr/bin/env python3
# This file is part of OpenCV project.
# It is subject to the license terms in the LICENSE file found in the top-level
# directory of this distribution and at http://opencv.org/license.html.

'''
Auto white balance using FC4: https://github.com/yuanming-hu/fc4

Color constancy is a method to make colors of objects render correctly on a photo.
White balance aims to make white objects appear white on an image and not a shade of any
other color, independent of the actual light setting. White balance correction creates
a neutral looking coloring of the objects, and generally makes colors look more similar
to their 'true' colors under different light conditions.

Given an RGB image, the FC4 model predicts scene illuminant (R,G,B). We then apply
the illuminant to the image, applying the correction in the linear RGB space.
The transformation between linear and sRGB spaces is done as described in the sRGB standard,
which is a nonlinear Gamma correction with exponent 2.4 and extra handling of very small values.
This sample is written for 8bit images. The FC4 model accepts RGB images with applied Gamma scaling.

The training of the FC4 model was done on the Gehler-Shi dataset. The dataset includes
568 images and ground truth corrections, as well as ground truth illuminants. The linear
RGB images from the dataset were used with Gamma correction of 2.2 applied.

The model is a pretrained fold 0 of a training pipeline on the Gehler-Shi dataset, from the PyTorch
implementation of the FC4 algorithm by Mateo Rizzo. The model was converted from a .pth file to onnx
using torch.onnx.export. The model can be downloaded in the following link:
https://raw.githubusercontent.com/MykhailoTrushch/opencv/d6ab21353a87e4c527e38e464384c7ee78e96e22/samples/dnn/models/fc4_fold_0.onnx

Copyright (c) 2017 Yuanming Hu, Baoyuan Wang, Stephen Lin
Copyright (c) 2021 Matteo Rizzo

Licensed under the MIT license.

References:

Yuanming Hu, Baoyuan Wang, and Stephen Lin. “FC⁴: Fully Convolutional Color
Constancy with Confidence-Weighted Pooling.” CVPR, 2017, pp. 4085–4094.

Implementations of FC4:
https://github.com/yuanming-hu/fc4/
https://github.com/matteo-rizzo/fc4-pytorch

Lilong Shi and Brian Funt, "Re-processed Version of the Gehler Color
Constancy Dataset of 568 Images," accessed from http://www.cs.sfu.ca/~colour/data/

“IEC 61966-2-1:1999 – Multimedia Systems and Equipment – Colour Measurement and Management –
Part 2-1: Colour Management – Default RGB Colour Space – sRGB.” IEC Standard, 1999.
'''

import argparse
import sys
import numpy as np
import cv2 as cv

from common import *


# Normalization constant for 8bit values
NORMALIZE_FACTOR = 1.0 / 255.0

# sRGB to linear conversion constants (or vice versa):
# SRGB_THRESHOLD / LINEAR_THRESHOLD: breakpoints between linear and gamma regions
# SRGB_SLOPE: slope of the linear segment near black
# SRGB_ALPHA: offset to ensure continuity at the threshold
# SRGB_EXP: gamma exponent
SRGB_THRESHOLD = 0.04045
SRGB_ALPHA     = 0.055
SRGB_SLOPE     = 12.92
SRGB_EXP       = 2.4
LINEAR_THRESHOLD = 0.0031308
EPS = 1e-10

def srgb_to_linear(rgb: np.ndarray) -> np.ndarray:
    low  = rgb / SRGB_SLOPE
    high = np.power((rgb + SRGB_ALPHA) / (1.0 + SRGB_ALPHA), SRGB_EXP, dtype=np.float32)
    return np.where(rgb <= SRGB_THRESHOLD, low, high).astype(np.float32)

def linear_to_srgb(lin: np.ndarray) -> np.ndarray:
    low  = lin * SRGB_SLOPE
    high = (1.0 + SRGB_ALPHA) * np.power(lin, 1.0 / SRGB_EXP, dtype=np.float32) - SRGB_ALPHA
    return np.where(lin <= LINEAR_THRESHOLD, low, high).astype(np.float32)

def correct(bgr8u: np.ndarray, illum_rgb_linear: np.ndarray) -> np.ndarray:
    assert bgr8u.dtype == np.uint8 and bgr8u.ndim == 3 and bgr8u.shape[2] == 3

    bgr = bgr8u.astype(np.float32) * NORMALIZE_FACTOR
    lin = srgb_to_linear(bgr)
    e_r = max(float(illum_rgb_linear[0]), EPS)
    e_g = max(float(illum_rgb_linear[1]), EPS)
    e_b = max(float(illum_rgb_linear[2]), EPS)
    s3 = np.float32(np.sqrt(3.0))
    corr_bgr = np.array([e_b * s3 + EPS,
                         e_g * s3 + EPS,
                         e_r * s3 + EPS],
                        dtype=np.float32)

    corrected = lin / corr_bgr.reshape(1, 1, 3)

    max_val = float(corrected.max()) + EPS
    corrected /= max_val
    corrected = np.clip(corrected, 0.0, 1.0)

    srgb = linear_to_srgb(corrected)

    out_bgr8 = (srgb * 255.0 + 0.5).astype(np.uint8)
    return out_bgr8

def annotate(img_bgr: np.ndarray, title: str) -> None:
    fs = max(0.5, min(img_bgr.shape[1], img_bgr.shape[0]) / 800.0)
    th = max(1, int(round(fs * 2)))
    cv.putText(img_bgr, title, (10, 30), cv.FONT_HERSHEY_SIMPLEX, fs, (0,255,0), th)

def get_args_parser(func_args):
    backends = ("default", "openvino", "opencv", "vkcom", "cuda", "webnn")
    targets = ("cpu", "opencl", "opencl_fp16", "ncs2_vpu", "hddl_vpu", "vulkan",
               "cuda", "cuda_fp16")

    p = argparse.ArgumentParser(add_help=False)
    p.add_argument('--zoo',
                   default=os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models.yml'),
                   help='An optional path to file with preprocessing parameters.')
    p.add_argument("--input", help="Path to input image", default="castle.png")
    p.add_argument('--backend', default="default", type=str, choices=backends,
            help="Choose one of computation backends: "
            "default: automatically (by default), "
            "openvino: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
            "opencv: OpenCV implementation, "
            "vkcom: VKCOM, "
            "cuda: CUDA, "
            "webnn: WebNN")
    p.add_argument('--target', default="cpu", type=str, choices=targets,
            help="Choose one of target computation devices: "
            "cpu: CPU target (by default), "
            "opencl: OpenCL, "
            "opencl_fp16: OpenCL fp16 (half-float precision), "
            "ncs2_vpu: NCS2 VPU, "
            "hddl_vpu: HDDL VPU, "
            "vulkan: Vulkan, "
            "cuda: CUDA, "
            "cuda_fp16: CUDA fp16 (half-float preprocess)")

    args, _ = p.parse_known_args()
    add_preproc_args(args.zoo, p, 'auto_white_balance', prefix="", alias="fc4")
    p = argparse.ArgumentParser(
        parents=[p],
        description="FC4 Color Constancy (ONNX): " \
        "predicts illuminant and applies white balance.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    return p.parse_args(func_args)



def main(func_args=None):
    args = get_args_parser(func_args)
    args.model = findModel(args.model, args.sha1)

    try:
        net = cv.dnn.readNetFromONNX(args.model)
        net.setPreferableBackend(get_backend_id(args.backend))
        net.setPreferableTarget(get_target_id(args.target))
    except cv.error as e:
        print(f"Error loading model: {e}", file=sys.stderr)
        sys.exit(1)

    img = cv.imread(findFile(args.input), cv.IMREAD_COLOR)
    if img is None:
        print(f"Cannot load image: {args.input}", file=sys.stderr)
        sys.exit(1)

    blob = cv.dnn.blobFromImage(
        img, scalefactor=args.scale, size=(img.shape[1], img.shape[0]),
        mean=args.mean, swapRB=args.rgb, crop=False, ddepth=cv.CV_32F
    )
    net.setInput(blob)

    try:
        out = net.forward()
    except cv.error as e:
        print(f"Forward error: {e}", file=sys.stderr)
        sys.exit(1)

    illum = out.astype(np.float32).reshape(-1)
    if out.size != 3:
        print("Error: model output of size not equal to 3 (should output 3 illuminants in RGB order)")
        sys.exit(-1)

    corrected = correct(img, illum)

    orig_vis = img.copy()
    corr_vis = corrected.copy()
    annotate(orig_vis, "Original")
    annotate(corr_vis, "FC4-corrected")
    stacked = np.hstack([orig_vis, corr_vis])
    cv.imshow("Original and Corrected Images", stacked)
    cv.waitKey(0)
    cv.destroyAllWindows()


if __name__ == "__main__":
    main()