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

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

Review comment (suggestion):
```
    net.setPreferableBackend(get_backend_id(args.backend))
    net.setPreferableTarget(get_target_id(args.target))
```

Code after:

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")