SWBNet: A Stable White Balance Network for sRGB Images (AAAI2023)

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

This paper is proposed to achieve stable white balance for the sRGB images with different color temperatures by learning the color temperature-insensitive features. paper link

Results presentation

Comparing with the state-of-the-art methods, our SWBNet has a stable and superior performance for the sRGB images with different color temperatures.

Framework

A. The CTIF extractor and CT-contrastive loss work together to learn the color temperature-insensitive features for achieving stable WB performance. B. The CTS-oriented transformer corrects multiple color temperature shifts differently to improve WB accuracy, especially for the multi-illumination sRGB images.

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_CTIF(l9el); Others are coming soon...
• 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 './example_images/1127_D.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...