WBFlow: Few-shot White Balance for sRGB Images via Reversible Neural Flows

1 Introduction

The color cast in sRGB images is caused by improper color temperature settings in the white balance (WB) module and then exacerbated by non-linear color renderings in image signal processor (ISP) [Afifi et al., 2019a]. The sRGB White balance (sRGB-WB) methods aim to correct such color cast without accessing raw values. It is an emerging direction and has an essential impact on some high-level computer vision tasks, such as object recognition and image segmentation [Afifi and Brown, 2019].

To perform accurate sRGB-WB, an ideal way is to reverse color-cast sRGB images to unprocessed raw values for linear correcting and then re-render them to obtain white-balanced sRGB images [Afifi and Brown, 2020a]. However, this process is challenging for sRGB-WB methods because the essential reversibility is difficult to achieve without raw values [Chakrabarti et al., 2014]. In fact, instead of the ideal WB process, recent sRGB-WB methods [Afifi et al., 2019a; Afifi et al., 2019b; Afifi and Brown, 2020b] directly utilize irreversible polynomial kernel functions [Hong et al., 2001] to train the mappings between color-cast sRGB images and ground truths (GTs). These trained mappings are not appropriate on untrained cameras due to the inherent color differences among sRGB images from different cameras [Afifi, 2021]. This seriously limits cross-camera generalization in real-world applications. In Figure 1(b), the representative sRGB-WB method, KNN-WB [Afifi et al., 2019a], results in a noticeable greenish cast on the untrained camera.

∗Corresponding Author.

Few endeavors have been made to address this problem. The only method, Deep WB [Afifi and Brown, 2020a], proposes modeling the ideal WB process end-to-end via U-Net [Ronneberger et al., 2015]. It further generalizes to different cameras by jointly training multi-camera sRGB images. However, due to the lossy structure of U-Net and the vanilla training strategy, Deep WB becomes irreversible and suppresses camera specificity. Thus, in Figure 1(c), although Deep WB improves generalization to the untrained camera, it has a degraded performance on the trained camera.

In this paper, we propose an sRGB-WB network called WBFlow that achieves superior performance on both trained and untrained cameras. Unlike existing irreversible sRGBWB methods, we take advantage of neural flows [Kingma and Dhariwal, 2018] to enable the reversibility of WBFlow. Specifically, we design Reversible Non-linear Rendering Transformation (RNRT) that simulates ISP color renderings by additive coupling layers [Kingma and Dhariwal, 2018]. Due to the inherent reversibility of additive coupling layers, RNRT can losslessly render color-cast sRGB images back to their pseudo-raw features via forward propagation. Then, to correct the color cast of pseudo-raw features, we present Reversible Linear Correction Transformation (RLCT) that simulates white balance and color transformation matrices as reversible 1 × 1 convolutions [Kingma and Dhariwal, 2018]. RLCT helps RNRT lossless back-propagation to generate white-balanced sRGB images by separating and correcting the color information from pseudo-raw features. Further, inspired by the fact that inter-camera transformation can be achieved by learning the mapping between raw values from different cameras [Bani´c et al., 2017], we generalize WBFlow to multiple cameras via inserting a camera transformation (CT) layer in pseudo-raw space. Specifically, by utilizing a few sRGB images from an untrained camera, we update the weightings of CT to enhance the camera specificity of pseudo-raw features, which cooperates with RNRT and RLCT to generate white-balanced sRGB images. Our contributions are as follows:

1. We propose WBFlow that generalizes well to multiple cameras via reversible neural flows and few-shot learning.
2. We propose RNRT and RLCT to simulate nonlinear color renderings and linear color transformations to enable the reversibility of WBFlow.
3. We propose CT that generalizes WBFlow to untrained cameras by updating its weights via few-shot learning.
4. Extensive experiments show that WBFlow achieves superior multi-camera generalization on three public datasets and our rendered multi-camera sRGB dataset.

2 RelatedWorks

White Balance Methods for Raw Images. The white balance methods for raw images fall into computational color constancy (CCC) research. Most CCC methods [Gijsenij and Gevers, 2007; Foster, 2011; Hu et al., 2017; Lo et al., 2021; Song et al., 2021; Xu et al., 2021; Tang et al., 2022; Zhang et al., 2022] achieve the goal of correcting the color cast in raw images by estimating illumination colors. They make up the WB module of camera ISPs. However, none of these methods can be applied to sRGB images [Afifi et al., 2019a], since the linear color cast of raw images is broken due to non-linear ISP renderings. Different from CCC methods, our WBFlow focuses on correcting the non-linear color cast of sRGB images.

White Balance Methods for sRGB Images. Recent sRGB-WB methods corrected the color cast in sRGB images by optimizing non-linear mappings (exemplar-based method)

[Afifi et al., 2019a; Afifi et al., 2019b; Afifi and Brown, 2020b; Afifi and Brown, 2019; Afifi et al., 2020] or modeling ideal WB process (DNN-based method) [Afifi and Brown, 2020a]. In detail, exemplar-based methods combined trained non-linear mappings to correct inputs. Since these mappings are irreversible, the specificities of trained cameras are mixed and applied to encode white-balanced sRGB images during testing, making them biased towards trained cameras rather than test ones. Inevitably, the performances of exemplarbased methods are inferior on untrained cameras (Figure 1(b)). To solve this problem, the DNN-based method, Deep WB, proposed rendering inputs back to pseudo-raw features by U-Net encoder, then correcting pseudo-raw features and re-rendering them to white-balanced sRGB images by U-Net decoder. It also allowed WB manipulation in sRGB images by setting two additional decoders to render inputs with 2850 Kelvin(K) and 7500K. However, due to lossy pooling layers, Deep WB is irreversible, making the content of pseudoraw features less complete than actual raw values and further restrictingWB accuracy. Further, the joint training strategy suppressed the specificities of different cameras [Huang et al., 2020], and hence together with irreversible structure, limited the improvement of multi-camera generalization (Figure 1(c)). Our WBFlow belongs to DNN-based methods and achieves superior generalization to multiple cameras via reversible neural flows and few-shot learning.

3 White Balance for sRGB Image via Reversible Flows and Few-shot Learning

In this section, we introduce WBFlow that aims to achieve superior generalization on multiple cameras. We start with formulating the ideal WB process for sRGB images. Then, we introduce details of WBFlow and how it generalizes to multiple cameras via few-shot learning.

3.1 Ideal White Balance Process for sRGB Images

An sRGB image Inct that is captured by camera n and rendered by color temperature ct can be formed as [Afifi et al., 2019a]:

where Inraw is a raw image from camera n, Wct is a WB module with color temperature ct, Tn is a color transformation matrix for camera n, fn(·) is a non-linear rendering function, including color enhancement, tone manipulation, and gamma encoding et al.. Corresponding correct sRGB image Inwb , i.e. GT, is obtained by manually setting Wct according to actual illumination color. Here, we denote this ideal setting as Wideal. Therefore, the relationship between Inwb and Inct can be modeled as:

From Equation 2, the ideal WB process is reversible to render color-cast sRGB image toward its white-balanced version while preserving complete content. Specifically, Inct is first reversed back to its color-cast linear values via f -1 n (·); then, Tn transforms color-cast linear values into raw values and then WidealW-1 ct white balance them; finally, fn(·) and Tn render white-balanced raw values to obtain Inwb .

As discussed in Section 1 and 2, existing sRGB-WB methods [Afifi et al., 2019a; Afifi et al., 2020; Afifi and Brown, 2020b; Afifi et al., 2021; Afifi and Brown, 2020a] are irreversible and thus fail to satisfy Equation 2, which limits their generalization to multiple cameras. In this paper, we propose an sRGB-WB network named WBFlow that satisfies the reversibility of Equation 2 and generalizes well to cameras via few-shot learning. Details will be introduced as follows.

3.2 White Balance for sRGB Images via Reversible Transformations

To achieve accurate white balance on multiple cameras, our primary goal is to appropriately model the ideal WB process. According to Equation 2, this process has two types of modules: nonlinear color rendering function (fn(·)) and linear color transformations (Wct, Wideal, Tn). Deep WB [Afifi and Brown, 2020a] models these two kinds of modules via identical irreversible structures, which only improves multicamera generalization to some extent. To solve this problem, different from Deep WB, we model two kinds of modules by individual reversible functions:

where F(·) is a reversible non-linear function that losslessly renders color-cast sRGB image Inct back to pseudo-raw features, M(·) is a reversible linear function that corrects the color cast of pseudo-raw features and preserves the content information, F -1 (·) is the inverse of F(·) that re-renders white-balanced pseudo-raw features to corresponding sRGB image ˆInwb . Theoretically, due to F(·) andM(·), Equation 3 satisfies the reversibility of Equation 2 to generate accurate white-balanced sRGB images on different cameras. Following are the implementation details.

Reversible Non-linear Rendering Transformation

As discussed, the reversible non-linear function F(·) aims to obtain pseudo-raw features with lossless content. Deep WB

[Afifi and Brown, 2020a] cannot satisfy it because of lossy pooling layers. Here, considering color renderings included in F(·) are independent [Ramanath et al., 2005], we model them by a sequence of reversible bijective sub-functions: F = F1 ◦ F2 ◦ · · · ◦ FK. This way, the relationship between color-cast sRGB image Inct and intermediate feature znct,k is:

where znct,K = znct is pseudo-raw feature.

In practice, we adopt additive coupling layers [Kingma and Dhariwal, 2018] to model {Fk}, which are effective in simulating nonlinear reversible transformations [Kingma and Dhariwal, 2018; An et al., 2021]. In our WBFlow, the forward computation of the additive coupling layer is:

where Split(·) splits znct,k into two parts along channel dimension to obtain zn,a ct,k and zn,b ct,k . Conv(·) renders colors of zn,a ct,k with 3 × 3 convolutions and relu functions. Concat(·) concatenates ˆzn,a ct,k and zn,b ct,k to form rendered feature znct,k+1 . The reverse computation of the additive coupling layer is easily derived from Equation 5. We can superimpose additive coupling layers to establish the reversible mapping between input color-cast sRGB image and pseudo-raw feature as Equation 4. However, as in Equation 5, the additive coupling layer leaves some channels of intermediate feature unchanged due to Split(·). To solve this problem, following [Kingma and Dhariwal, 2018], we use reversible 1×1 convolution that integrates all channels by fixing the channel numbers of input and output as the same. We also adopt Actnorm [Kingma and Dhariwal, 2018] to normalize features, an alternative to BatchNorm that accelerates the training times.

The proposed RNRT is shown in Figure 2, consisting of additive coupling layers, reversible 1 × 1 convolutions, and Actnorms. All components of RNRT are reversible, making content lossless during in forward and backward propagations. As shown in Figure 3(c), we can obtain pseudo-raw features with lossless content from color-cast input in Figure 3(a) via forward RNRT. Simultaneously, in Figure 3(b), the input can be restored accurately from pseudo-raw features in Figure 3(c) via reverse RNRT.

Reversible Linear Correction Transformation

In Equation 2 and 3, the reversible linear function M(·) is linear and reversible due to the WB module and color transformation matrix. To properly implement it while preserving lossless content information, inspired by [Huang and Belongie, 2017], we propose RLCT that separates color and content information from pseudo-raw features and corrects the color information by reversible 1 × 1 convolutions [Kingma and Dhariwal, 2018]. The formation is:

where channel mean and variance, µz and σz, represents the color information of pseudo-raw feature znct,k [Stricker and Orengo, 1995]. They also help to keep content information by normalizing znct , i.e., (znct - µz)/σz. For individual color information, we correct it by respectively forwarding µz and σz into linear mappings Tσ(·) and Tµ(·), both of which consist of three reversible 1×1 convolutions [Kingma and Dhariwal, 2018]. This way, in Figure 3(f), the pseudo-raw features are white-balanced by re-scaling lossless content with learned channel mean Tσ(σz) and variance Tµ(µz). Further, in Figure 3(e), the white-balanced sRGB image is successfully generated from white-balanced pseudo-raw features in Figure 3(f) via reverse propagation of RNRT. It has an almost identical color histogram as GT in Figure 3(d).

As introduced above, RNRT and RLCT are reversible and thus cooperate to ensure complete reversibility, which satisfies the requirement of Equation 2. Thus, our WBFlow can significantly improve the floor of WB accuracy on multiple cameras. As shown in Figure 4(a), with the same multicamera training strategy of Deep WB, our WBFlow converges with a much lower l1 loss on multiple cameras.

3.3 Camera Generalization via Few-shot Learning

According to Equation 2, the actual raw features should be specific to various cameras. However, existing methods suppressed such specificities due to combined trained mappings [Afifi et al., 2019a; Afifi and Brown, 2020b] or joint multicamera training strategy [Afifi and Brown, 2020a]. Inevitably, this will result in a limited multi-camera generalization. To enhance the camera specificity of pseudo-raw feature, we introduce few-shot learning into WBFlow, which is based on the fact that inter-camera transformation can be achieved by learning the mapping between raw color values from two cameras [Gao et al., 2017]. In Figure 2, we set the camera transformation named CT after RNRT. For camera n, Equation 6 can be re-written as:

where ∗ is the group convolution operator, wn cam is a set of weights learned from the camera labels by two reversible neural flow combinations. This strategy ensures that the learned weights are unique to the different cameras. To adapt WBFlow to camera n, we fix the pre-trained parameter of RNRT, then train the new weights in Equation 7 by minimizing the l1 loss with a few samples. This way, the specificity of pseudo-raw features to camera n is enhanced, thus improving the generalization effect. In Figure 4(b), the ∆E2000 values of WBFlow on the new camera, PanasonicGX1, are always lower than these of Deep WB, regardless of the number of training cameras. More details of few-shot learning settings will be introduced in Section 4.

4 Experiments

In this section, we compare the effectiveness of WBFlow with state-of-the-art sRGB-WB methods in terms of multi-camera generalization, inference time, and few-shot learning. In addition, ablation experiments are designed to evaluate the effectiveness of each part in WBFlow.

4.1 Datasets

Public Datasets. Following [Afifi and Brown, 2020a], we randomly selected 12000 sRGB images from the first fold of Set1 [Afifi et al., 2019a; Cheng et al., 2014] to train WBFlow.

We used the remaining two folds (Set1-Test), Set2 [Afifi et al., 2019a] and rendered cube dataset [Afifi et al., 2019a; Bani´c et al., 2017] for evaluation.

Rendered Multi-camera sRGB Dataset. We collected a multi-camera sRGB dataset to evaluate the multi-camera generalization effect. Specifically, we selected and compiled 184 groups of raw images from the NUS dataset [Cheng et al., 2014]. In each group, the raw images are consistent in the scenes and different in the cameras: Canon1DsMkIII, Canon600D, FujifilmXM1, NikonD5200, OlympusEPL6, PanasonicGX1, SamsungNX2000, and SonyA57. To obtain color-cast sRGB versions of these images, following [Afifi et al., 2019a], we rendered them by Adobe Camera Raw with five common color temperatures (2850 K, 3800 K, 5500 K, 6500 K, and 7500 K) and camera standard photo finishing. We obtain GTs by manually selecting the correct color temperature from the middle gray patches in the color checker of each raw image. The rest of the operations remain unchanged. Our multi-camera sRGB dataset contains 7360 sRGB images with 184 scenes, five color temperatures, and eight cameras.

4.2 Implementation Details

Loss Function. Following [Afifi and Brown, 2020a], we apply l1 loss to trainWBFlow:

Training and Testing Detail. We implemented WBFlow on Pytorch with CUDA support and used the Adam [Kingma and Ba, 2014] with β1 = 0.9 and learning rate 10-4 to optimize it. For the experiments with all training images, we trained WBFlow for 340000 iterations with batch size 4. While for few-shot experiments, we trained CT for 15000 iterations. We used color jittering, average blur, geometric rotation, and flipping to augment data. During testing, following [Afifi and Brown, 2020a], we resized all input images to a maximum dimension of 656 pixels and set a color mapping procedure to compute the final white-balanced sRGB images.

Error Metric. We used the same error metrics as existing sRGB-WB methods [Afifi and Brown, 2020a; Afifi et al., 2019a]: mean angle error (MAE) and ∆E2000 [Sharma et al., 2005]. We reported the mean, first quantile (Q1), second quantile (Q2), and third quantile (Q3) for evaluation. A lower error metric denotes a better sRGB-WB performance.

4.3 Multi-camera Generalization Evaluation

Table 1 shows the quantitative results of our WBFlow and state-of-the-art sRGB-WB methods on Set1-Test (trained cameras), Set2 (untrained cameras), and rendered cube dataset (untrained cameras). The results of MixedWB on Set1-Test and Set2 are computed by its code, and others are collected from [Afifi et al., 2021] and [Afifi and Brown, 2020a].

Quantitate Comparison. From Table 1, due to the irreversible and biased mapping, the exemplar-based methods, KNN-WB and MixedWB, generalized badly to untrained cameras. For example, their mean values of ∆E2000 on Set2 and rendered cube dataset are 56.42% and 36.97% higher than those in Set1-Test. Deep WB alleviated this by modeling the ideal WB process by U-Net. However, since UNet is irreversible, Deep WB performed inferiorly on trained cameras. That is, although Deep WB reduced the performance of KNN-WB in the mean of ∆E2000 about 12.50% and 19.19% in Set2 and rendered cube dataset, its mean of ∆E2000 is 5.31% worse than KNN-WB in Set1-test. In contrast, in Set1-test, our WBFlow outperformed the KNN-WB (ranking second) and the Deep WB (most related method) by about 12.57% and 16.98% in the mean of ∆E2000, respectively. Simultaneously, in Set2 and rendered cube dataset, our WBFlow still outperformed all methods in most metrics, e.g., it greatly improved the accuracy of Deep WB in the mean of ∆E2000 and MAE about 5.31% and 6.40%. Similar results appear in Set2. These results verify the effectiveness of our WBFlow in multi-camera generalization.

Inference Time Comparison. From Table 1, our WBFlow achieves superior WB performance on all datasets with almost the same inference time as Deep WB (the most related method). Compared with KNN-WB with the fastest inference time, our WBFlow outperforms it by about 12.57%, 17.14% and 24.65% for the mean of ∆E2000 in three datasets.

Qualitative Comparison. To qualitatively compare the multi-camera generalization, we show the generated whitebalanced sRGB images and their color histograms of KNNWB, Deep WB, and WBFlow on eight cameras in Figure 5. From it, WBFlow achieves the most similar colors as GTs on all cameras, consistent with the quantitate comparison.

4.4 Few Shot Evaluation

To validate few-shot capacity, we compared the performances of WBFlow with the most related method, Deep WB[Afifi and Brown, 2020a], on four untrained cameras. We retrained Deep WB and WBFlow by randomly selecting 5000 sRGB images from the first fold of Set1 and used four untrained cameras in the last 84 groups of our rendered multi-camera sRGB dataset for evaluation: OlympusEPL6, PanasonicGX1, SamsungNX2000, and SonyA57. Training details are followed by [Afifi and Brown, 2020a]. For few-shot experiments, we randomly selected K images (0, 5, 10, 20) for untrained cameras from the first 100 groups of our rendered multi-camera sRGB dataset to train the parameters of CT, which are 5.82% of our WBFlow's. To avoid randomness and disturbance, we repeated experiments 1000 times with randomly selected images and then computed the ∆E2000 average in Table 2.

From Table 2, WBFlow significantly outperforms Deep WB even with no shots (K=0) due to the reversible structure. For example, the average ∆E2000 values of WBFlow are smaller than these of Deep WB by about 17.44%, 23.09%, 21.68%, and 23.45% in Q2 on four untrained cameras. This superiority was enhanced when we updated CT with K=5, 10, 20 to enhance the camera specificity. Especially, the fewshot learning effectively improves Q3 of average ∆E2000 on four untrained cameras. This indicates that our WBFlow with few-shot learning is superior in handling untrained cameras.

4.5 Ablation Analysis

To verify the effectiveness of each part in WBFlow, we conducted an ablation analysis on our rendered multi-camera sRGB dataset. Since Deep WB is the most related method, we report the average ∆E2000 of Deep WB and three variants of our WBFlow in Figure 6(a). Training and testing settings are the same as few-shot experiments. Further, we compared the influences of flow numbers for RNRT in Figure 6(b) and the camera specificity verification for CT in Figure 7.

Variants Comparison. From Figure 6(a), the average ∆E2000 values of RNRT on eight cameras are much worse than Deep WB, which shows that only modeling color renderings hardly improves the multi-cameras generalization. This phenomenon is significantly mitigated when we integrate RNRT and RLCT. Specifically, compared with only RNRT, RNRT&RLCT work together to reduce the average ∆E2000 values by about 50% on eight cameras. RNRT&RLCT also considerably outperforms Deep WB by 13% on eight cameras. This is because both RNRT and RLCT are reversible to effectively model the ideal WB process. When we add CT, the multi-camera generalization is further improved, as in the smallest average ∆E2000 of RNRT&RLCT&CT on eight cameras.

Flows Numbers Influence for RNRT. To explore the optimal flow numbers of RNRT, in Figure 6(b), we compared the average ∆E2000 of RNRT with 4, 8, 12, and 16 flows on eight cameras. It can be seen that the complexity of the RNRT becomes greater as the number of streams increases to better simulate ISP nonlinear renderings. The average ∆E2000 of RNRT is minimized at a number of 16, which means that 16 is the optimal flow number.

Camera Specificity Verification. Since channel average and variance can represent the color information of pseudoraw features [Stricker and Orengo, 1995], we compute them for white-balanced pseudo-raw features without and with CT on eight cameras to verify the camera specificity in Figure 7. With CT, the difference between mean/variance for all white balance pseudo-raw feature channels becomes significantly larger on eight cameras. This verifies that CT can significantly enhance the camera specificity of pseudo-raw features, thus improving multi-camera generalization.

5 Conclusion

In this paper, we propose an sRGB-WB network named WBFlow, which performs superior white balance for sRGB images and generalizes well to multiple cameras. Unlike existing irreversible sRGB-WB methods that fail to model the ideal WB process, WBFlow successfully models this process through reversible RNRT and RLCT. Furthermore, we generalize WBFlow to multiple cameras by enhancing the camera characteristic of pseudo-raw features via few-shot learning. Extensive experiments indicate the superiority of our WBFlow in multi-camera generalization and few-shot sRGBWB tasks. Ablation analysis shows the effectiveness of each part of WBFlow and their optimal combinations.

Acknowledgements

This work was supported by the national key R & D program intergovernmental international science and technology innovation cooperation project (2021YFE0101600).