WB LUTs: Contrastive Learning for White Balancing Lookup Tables

TL;DR

This work introduces WB LUTs, a 3D LUT-based framework for high-resolution, real-time white balance correction that bypasses bottlenecks in post-capture upsampling. By integrating a scene classifier with a contrastive learning framework and a novel hard sample mining strategy, the method learns illumination-aware, scene-agnostic representations to optimize color correction. The approach achieves near-state-of-the-art WB quality on benchmark datasets while offering 300x faster inference and 12.7x lower memory than competing models, and it emphasizes LAB color space with hard positives for best performance. The results demonstrate strong quantitative gains and qualitative improvements in high-resolution WB outputs, with potential for rapid deployment in ISP pipelines; future work notes local color correction for scenes with multiple illuminants.

Abstract

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.

Figures (6)

• Figure 1: Comparison of performance and WB correction quality of our method with the state-of-the-art, on the Rendered Cube dataset banic2017unsupervised. Performance is gauged both by the run time, on the vertical axis, and the model memory, which is proportional to the size of the marker for each model.
• Figure 2: Downsampled input $I_{\downarrow}$ is fed to the scene classifier to generate weights for basis LUT fusion. The adaptive LUT is used to correct the high resolution image.
• Figure 3: Two images, both rendered with the Cloudy WB setting, illustrate how the same white balance rendering can result in different color mappings for different scenes.
• Figure 4: Overview of proposed contrastive learning approach for LUTs. For an image $I_A$ and its ground truth $I^\text{gt}_A$, we select an image of the same scene rendered with a different color temperature as negative sample $I^N_A$. Ground truth image of a different scene $I^\text{gt}_B$ is color mapped to generate a hard positive sample $I^{P}_A := M^\star\psi(I^\text{gt}_B)$. These hard positive and negative samples are used in the contrastive learning framework supervised with triplet loss.
• Figure 5: Results from base 3D LUTs vs WB LUTs, trained with the proposed contrastive learning framework, with $\Delta E 2000$ for each output overlaid onto the image.