WB LUTs: Contrastive Learning for White Balancing Lookup Tables
Sai Kumar Reddy Manne Michael Wan
Automatic white balancing (AWB), one of the first steps in an integrated
signal processing (ISP) pipeline, aims to correct the color cast induced by the
scene illuminant. An incorrect white balance (WB) setting or AWB failure can
lead to an undesired blue or red tint in the rendered sRGB image. To address
this, recent methods pose the post-capture WB correction problem as an
image-to-image translation task and train deep neural networks to learn the
necessary color adjustments at a lower resolution. These low resolution outputs
are post-processed to generate high resolution WB corrected images, forming a
bottleneck in the end-to-end run time. In this paper we present a 3D Lookup
Table (LUT) based WB correction model called WB LUTs that can generate high
resolution outputs in real time. We introduce a contrastive learning framework
with a novel hard sample mining strategy, which improves the WB correction
quality of baseline 3D LUTs by 25.5%. Experimental results demonstrate that the
proposed WB LUTs perform competitively against state-of-the-art models on two
benchmark datasets while being 300 times faster using 12.7 times less memory.
Our model and code are available at https://github.com/skrmanne/3DLUT_sRGB_WB.