mahmoudnafifi / C5

Reference code for the paper "Cross-Camera Convolutional Color Constancy" (ICCV 2021)
Apache License 2.0
105 stars 19 forks source link
color-constancy color-correction colorconstancy computational-photography cross-camera deep-learning hypernetworks iccv2021 illuminant-estimation machine-learning pytorch white-balance whitebalance

Cross-Camera Convolutional Color Constancy, ICCV 2021 (Oral)

Mahmoud Afifi1,2, Jonathan T. Barron2, Chloe LeGendre2, Yun-Ta Tsai2, and Francois Bleibel2

1York University   2Google Research

Paper | Poster | PPT | Video

C5_teaser

Reference code for the paper Cross-Camera Convolutional Color Constancy. Mahmoud Afifi, Jonathan T. Barron, Chloe LeGendre, Yun-Ta Tsai, and Francois Bleibel. In ICCV, 2021. If you use this code, please cite our paper:

@InProceedings{C5,
  title={Cross-Camera Convolutional Color Constancy},
  author={Afifi, Mahmoud and Barron, Jonathan T and LeGendre, Chloe and Tsai, Yun-Ta and Bleibel, Francois},
  booktitle = {The IEEE International Conference on Computer Vision (ICCV)},
  year={2021}
}

C5_figure

Code

Prerequisite

Training

To train C5, training/validation data should have the following formatting:

- train_folder/
       | image1_sensorname_camera1.png
       | image1_sensorname_camera1_metadata.json
       | image2_sensorname_camera1.png
       | image2_sensorname_camera1_metadata.json
       ...
       | image1_sensorname_camera2.png
       | image1_sensorname_camera2_metadata.json
       ...

In src/ops.py, the function add_camera_name(dataset_dir) can be used to rename image filenames and corresponding ground-truth JSON files. Each JSON file should include a key named either illuminant_color_raw or gt_ill that has the ground-truth illuminant color of the corresponding image.

The training code is given in train.py. The following parameters are required to set model configuration and training data information.

The following parameters are useful to control training settings and hyperparameters:

Testing

To test a pre-trained C5 model, testing data should have the following formatting:

- test_folder/
       | image1_sensorname_camera1.png
       | image1_sensorname_camera1_metadata.json
       | image2_sensorname_camera1.png
       | image2_sensorname_camera1_metadata.json
       ...
       | image1_sensorname_camera2.png
       | image1_sensorname_camera2_metadata.json
       ...

The testing code is given in test.py. The following parameters are required to set model configuration and testing data information.

In the images directory, there are few examples captured by Mobile Sony IMX135 from the INTEL-TAU dataset. To white balance these raw images, as shown in the figure below, using a C5 model (trained on DSLR cameras from NUS and Gehler-Shi datasets), use the following command:

python test.py --testing-dir-in ./images --white-balance True --model-name C5_m_7_h_64

c5_examples

To test with the gain multiplie, use the following command:

python test.py --testing-dir-in ./images --white-balance True --g-multiplier True --model-name C5_m_7_h_64_w_G

Note that in testing, C5 does not require any metadata. The testing code only uses JSON files to load ground-truth illumination for comparisons with our estimated values.

Data augmentation

The raw-to-raw augmentation functions are provided in src/aug_ops.opy. Call the set_sampling_params function to set sampling parameters (e.g., excluding certain camera/dataset from the soruce set, determine the number of augmented images, etc.). Then, call the map_raw_images function to generate a new augmentation set with the determined parameters. The function map_raw_images takes four arguments: