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

Code before:
#!/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
#
# 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)

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)

backends = (cv.dnn.DNN_BACKEND_DEFAULT, cv.dnn.DNN_BACKEND_HALIDE, cv.dnn.DNN_BACKEND_INFERENCE_ENGINE, cv.dnn.DNN_BACKEND_OPENCV,
            cv.dnn.DNN_BACKEND_VKCOM, cv.dnn.DNN_BACKEND_CUDA)
targets = (cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_OPENCL, cv.dnn.DNN_TARGET_OPENCL_FP16, cv.dnn.DNN_TARGET_MYRIAD, cv.dnn.DNN_TARGET_HDDL,
           cv.dnn.DNN_TARGET_VULKAN, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16)


def main():
    p = argparse.ArgumentParser(
        description="FC4 Color Constancy (ONNX): " \
        "predicts illuminant and applies Von Kries white balance."
    )
    p.add_argument("--model", required=True, help="Path to ONNX model file")
    p.add_argument("--input", required=True, help="Path to input image")
    p.add_argument("--scale", type=float, default=1.0, help="Input scaling factor (e.g., 1/255)")
    p.add_argument("--mean", type=float, nargs=3, default=[0.0, 0.0, 0.0],
                   help="Mean to subtract (B G R)")
    p.add_argument("--rgb", action="store_true", help="Swap BGR->RGB for model input")
    p.add_argument('--backend', choices=backends, default=cv.dnn.DNN_BACKEND_DEFAULT, type=int,
                    help="Choose one of computation backends: "
                         "%d: automatically (by default), "
                         "%d: Halide language (http://halide-lang.org/), "
                         "%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
                         "%d: OpenCV implementation, "
                         "%d: VKCOM, "
                         "%d: CUDA"% backends)
    p.add_argument('--target', choices=targets, default=cv.dnn.DNN_TARGET_CPU, type=int,
                        help='Choose one of target computation devices: '
                            '%d: CPU target (by default), '
                            '%d: OpenCL, '
                            '%d: OpenCL fp16 (half-float precision), '
                            '%d: NCS2 VPU, '
                            '%d: HDDL VPU, '
                            '%d: Vulkan, '
                            '%d: CUDA, '
                            '%d: CUDA fp16 (half-float preprocess)'% targets)
    args = p.parse_args()

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

    img = cv.imread(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=tuple(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 = extract_illuminant(out)
    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()


Review comment (suggestion):
You can directly use `mean=args.mean` instead of `mean=tuple(args.mean)`

Code after:
    )
    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()