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

Code before:

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, '

Review comment (suggestion):
```
    backends = ("default", "openvino", "opencv", "vkcom", "cuda")
    targets = ("cpu", "opencl", "opencl_fp16", "ncs2_vpu", "hddl_vpu", "vulkan", "cuda", "cuda_fp16")
```
Use strings

Code after:

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