vendor: OpenCV 5.0.0 snapshot at 40738fb16ceddb5fb3fea747585f7ce6abb0605b
This commit is contained in:
Gitea Mirror Bot
@@ -0,0 +1,202 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user