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

• Authors: Chunxiao Li, Xuejing Kang, Anlong Ming*
• Affiliation: School of Computer Science (National Pilot Software Engineering School), Beijing University of Posts and Telecommunications

This paper is proposed to build a model that not only performs superior white balance for sRGB images but also generalizes to multiple cameras well.. paper link

Framework

We propose the Reversible Non-linear Rendering Transformation and Reversible Linear Correction Transformation to ensure the reversibility of WBFlow, which significantly improves the floor of WB accuracy for multiple cameras. The Camera Transformation is then applied in the pseudo-raw space to generalize WBFlow to multiple cameras via few-shot learning.

Results presentation

Comparing with the state-of-the-art methods, our WBFlow has a stable and superior performance for the sRGB images from different cameras.

Rendered Multi-camera sRGB Dataset

Intro

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. In each group, the raw images are consistent in the scenes and differ in the cameras: Canon1DsMkIII, Canon600D, FujifilmXM1, NikonD5200, OlympusEPL6, PanasonicGX1, SamsungNX2000, and SonyA57. To obtain the color-cast sRGB versions of these images, we use the Adobe Camera Raw in Photoshop to render them with five common color temperatures (2850 K, 3800 K, 5500 K, 6500 K, and 7500 K) and camera standard photo finishing. We obtain the corresponding 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. In total, our multi-camera sRGB dataset contains 7360 sRGB images with 184 scenes, five color temperatures, and eight cameras.

Dataset is available: Data

The folds are coming soon...

Experiment

Requirements

• Python 3.8.3
• pytorch (1.8.0)
• torchvision (0.8.1)
• tensorboard (optional)
• numpy
• Pillow
• tqdm
• matplotlib
• scipy
• scikit-learn

Testing

• Pretrained models: Net
• Please download them and put them into the floder ./model/

Testing single image

• To test single image, changing '--input' in demo.sh and run it. The result is save in the folder 'result_images'.

demo.sh
python demo_single_image.py --input '**.JPG' --output_dir './result_images'

Testing multiple images

• Public datasets are available: Rendered WB dataset (Set1, Set2, Cube)
• To test multiple images, changing '--input_dir', '--gt_dir' and '--output_dir' in demo_images.py and run it.

python demo_images.py --input_dir --gt_dir --output_dir

Training

• Training fold is formed according to Deep White-balance Editing (CVPR 2020)
• Training fold is available: Training Fold
• Training data can be loaded from: Rendered WB dataset-Set1
• To train the model, changing '--training_dir', '--data-name' and '--test-name' in train.py and run it.

python train.py --training_dir --data-name --test-name

Citation

• Paper is available: paper link
• Citing format is coming soon...