Repository: opencv/opencv
File: samples/dnn/auto_white_balance.py (Python)

Code before:
# 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
#
# 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.
#
# Yuanming Hu, Baoyuan Wang, and Stephen Lin. “FC⁴: Fully Convolutional Color
# Constancy with Confidence-Weighted Pooling.” CVPR, 2017, pp. 4085–4094.

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


def extract_illuminant(out: np.ndarray) -> np.ndarray:
    flat = out.astype(np.float32).reshape(-1)
    assert flat.size >= 3, "ONNX output has fewer than 3 values"
    return flat[:3]

def srgb_to_linear(rgb: np.ndarray) -> np.ndarray:
    a = 0.055
    low  = rgb / 12.92
    high = np.power((rgb + a) / (1.0 + a), 2.4)
    return np.where(rgb <= 0.04045, low, high).astype(np.float32)

def linear_to_srgb(lin: np.ndarray) -> np.ndarray:
    a = 0.055
    low  = lin * 12.92
    high = (1.0 + a) * np.power(lin, 1.0/2.4) - a
    return np.where(lin <= 0.0031308, 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
    rgb = cv.cvtColor(bgr8u, cv.COLOR_BGR2RGB).astype(np.float32) / 255.0
    lin = srgb_to_linear(rgb)

    s3 = np.sqrt(3.0).astype(np.float32)
    corr = (illum_rgb_linear.astype(np.float32) * s3) + 1e-10  # [R,G,B]
    corrected = lin / corr.reshape(1, 1, 3)

    max_val = corrected.max() + 1e-10
    normalized = corrected / max_val

    srgb = np.clip(linear_to_srgb(normalized), 0.0, 1.0)
    out_rgb8 = (srgb * 255.0 + 0.5).astype(np.uint8)
    return cv.cvtColor(out_rgb8, cv.COLOR_RGB2BGR)


Review comment (suggestion):
You can specify these constant values globally and then use that, also add a small comment about what does these constants represent for better understanding.

Code after:
# 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