Time-Aware Auto White Balance in Mobile Photography (ICCV 2025)

This repository contains the code for the paper:

Time-Aware Auto White Balance in Mobile Photography

Mahmoud Afifi*, Luxi Zhao*, Abhijith Punnappurath, Mohamed A. Abdelsalam, Ran Zhang, and Michael S. Brown

AI Center-Toronto, Samsung Electronics

* Equal contribution

Figure: The time of day influences scene illumination, making it a valuable cue for improving illuminant estimation. Shown are white-balanced results (gamma-corrected for visualization) using our method with (1) colors only, (2) colors plus contextual (timestamp and geolocation) and capture data, and (3) ground truth (from a color chart). Angular errors show improvements when time-capture information is used.

📖 Overview

Cameras rely on auto white balance (AWB) to correct undesirable color casts caused by scene illumination and the camera's spectral sensitivity. This is typically achieved using an illuminant estimator that determines the global color cast solely from the color information in the camera's raw sensor image. Mobile devices provide valuable additional metadata—such as capture timestamp and geolocation—that offers strong contextual clues to help narrow down the possible illumination solutions. In this work, we propose a lightweight AWB model (~5K parameters) that processes two distinct inputs: (1) a time-capture feature, which combines contextual and capture information available on mobile devices and accessible via their camera ISPs, and (2) a histogram feature, which represents the image’s R/G and B/G chromaticity distributions. To validate our method, we introduce a dataset of 3,224 smartphone images captured at various times of day and under diverse lighting conditions. Each image is annotated with both ground-truth illuminants, determined using a color chart, and user-preferred illuminants, validated through a user study. This dataset provides a comprehensive benchmark for evaluating AWB methods. Our model achieves strong results, matching or outperforming significantly larger models.

Figure: The time of day provides valuable cues about likely illuminant colors in outdoor scenes. The figure on the left shows the correlated color temperature (CCT) of illuminants across different times of day and night. Excluding artificial lighting, images captured at noon, for example, exhibit different CCT ranges compared to those taken during sunrise or sunset. This observation motivated our lightweight model design (shown on the right), which leverages time-of-day cues. A convolutional network processes histogram features and combines them with time-capture features, followed by a lightweight MLP that estimates the scene’s illuminant chromaticity, which is then converted to a normalized RGB color.

🛠 Installation

You can set up the environment using either Conda or venv:

Option 1: Using Conda

# Create and activate a new conda environment
conda create -n time_aware_awb_env python=3.8 -y
conda activate time_aware_awb_env

# Install required packages
pip install -r requirements.txt

Option 2: Using venv

# Create and activate a virtual environment
python -m venv time_aware_awb_env
source time_aware_awb_env/bin/activate  # On Windows, use `.\time_aware_awb_env\Scripts\activate`

# Install required packages
pip install -r requirements.txt

To enable GPU acceleration, you must install a CUDA-compatible version of PyTorch. Please follow the official PyTorch installation instructions to install the correct version for your system.

🏞️ S24 Raw-sRGB Dataset

We present a new large-scale dataset of raw images, each paired with multiple sRGB ground-truth renditions and a denoised raw version. Each image is annotated with:

• A ground-truth illuminant, measured using a color chart placed in the scene during capture.
• A user-preference white-balance ground truth, selected to reflect real-world visual preferences and enhance scene aesthetics.

Figure: For each scene in our dataset, an image with a color chart placed in the scene was captured. The gray patches on the color chart were used to measure the ground-truth illuminant. These color-chart images were discarded, and only the images of the scenes without the color chart were used for the training, validation, and testing sets.

The dataset also includes additional annotations, such as scene type and lighting class. For more details, see here.

🗂️ Directory Structure

Your dataset directory should have the following structure:

/path/to/dataset/
│── train/
│   ├── raw_images/         # 16-bit PNG raw images (384×256, masked)
│   ├── data/               # JSON files with metadata and ground-truth illuminants
│── val/
│   ├── raw_images/         # 16-bit PNG raw images (384×256, masked)
│   ├── raw_images_wo_mask/ # 16-bit PNG raw images (384×256, unmasked)
│   ├── data/               # JSON files with metadata and ground-truth illuminants
│── test/
│   ├── raw_images/         # 16-bit PNG raw images (384×256, masked)
│   ├── raw_images_wo_mask/ # 16-bit PNG raw images (384×256, unmasked)
│   ├── data/               # JSON files with metadata and ground-truth illuminants

To generate this expected dataset structure, first download the S24 Raw-sRGB Dataset, then run the following command:

python dataset_preprocessing.py --dataset_path /path/to/s24/dataset/ --output_path /path/to/processed/dataset/

To change the names of masked and unmasked subfolders, refer to constants.py.

⚙️ Training

To train the model, use the following command:

python training.py --dataset_path /path/to/processed/dataset/ --capture_data iso shutter_speed flash time noise_stats snr_stats

Additional Arguments:

• --user_preference_gt: Train on user-preference-based ground truth.
• --hist_bins N: Change the histogram bin count (default: 48, e.g., --hist_bins 64).
• --exp_name name: Set a specific experiment name (default: "awb").
• --capture_data: Select one or more from iso, shutter_speed, flash, time, noise_stats, snr_stats.

Training typically takes about 10 minutes. We recommend naming models trained with user preference ground truth using the -pref suffix (e.g., awb-pref). After testing, you can report results using evaluation.py by specifying the base experiment name (e.g., awb), and the script will automatically display results for both the neutral and preference models. More details are provided below.

📊 Testing and Evaluation

To test the model, use:

python testing.py --dataset_path /path/to/processed/dataset/ --set_name test --output_path path/to/save/results

Additional Arguments:

• --exp_name name: Match the experiment name used during training (default: "awb").
• --without_mask: Test without applying masks.
• --set_name: Set name; options are "test" and "val".

The script saves estimated illuminant colors under --output_path, in a subfolder named with_mask or without_mask based on usage. If the model was trained using user-preference white balance, results will be saved in a pref_illum subfolder; otherwise, in gt_illum.

To compute the angular error between predicted and ground-truth illuminants, run:

python evaluation.py --dataset_path /path/to/processed/dataset/ --set_name test --results_dir /path/where/the/results/were/saved --exp_name awb

Additional Arguments:

• --latex_format: Format the output for inclusion in LaTeX tables.
• --exp_name name: Should match the experiment name used during testing. If you have two models with the same basename—one for neutral AWB and one for preference (e.g., awb and awb-pref)—and both have been tested, you can simply specify the basename (e.g., awb). The results for both models (neutral and preference AWB) will then be reported.
• --without_mask: Report results on the unmasked image set.

💻 Demo

To run our method on new images outside the dataset (the demo performs white balance using our model and renders the image to sRGB without advanced photofinishing):

python demo.py --dng_file /path/to/dng/img --jpg_file path/to/jpeg/img

Additional Arguments:

• --json_file: Use this if the JPEG file is unavailable. It should contain processed metadata extracted from the JPEG, similar to the JSON files in the data directory of our dataset.
• --output_dir: Directory to save the output image (default: .).
• --device: Device to run the model on: "gpu" or "cpu" (default: "gpu").
• --user_preference_gt: Use the user-preference trained model.
• --snr_stats: Use SNR stats during inference.

Note: The demo runs the model trained using time feature and ISO, shutter speed, and flash capture data. It assumes the input DNG is from the Samsung S24 Ultra main camera (used during training). For generalization to other cameras, refer to Section B.6 of our paper. The JPEG image is required to extract time/location metadata. The demo does not support noise stats, as we used Adobe Lightroom for denoising. However, a custom denoiser trained on the S24 noisy/clean raw pairs (provided in our dataset) can be used to compute noise stats and enable the use of models that rely on them. Also, note that the provided models without noise and SNR stats were trained later and may yield different results from those reported in the paper.

🏷️ Labeling Tool

Figure: We developed a GUI tool to streamline the annotation process. It includes navigation, white balance (WB) tools, scene/lighting classification, visualization, and dataset management. Annotators can select and copy a neutral WB from a color chart, interpolate between camera and neutral WB, and adjust the correlated color temperature (CCT) for cooler or warmer tones.

As described in Section C.2 of our paper, we developed a MATLAB graphical user interface (GUI) tool to facilitate the annotation process for our dataset.
The source code is available in the labeling_tool/ directory.

To use the labeling tool to annotate a new set of images, you must first pre-process the DNG and JPEG images captured by the camera. Specifically, run the pre_labeling.py script, which requires the following arguments:

• --dataset_path: to specify the path to the directory containing the DNG and JPG images extracted from the camera (both files should be located in the same directory and share the same filename, differing only by extension).

• --output_path: to define the output directory where the processed files will be saved.

After running the pre_labeling.py script, the following subdirectories will be created inside --output_path:

• raw_images: Processed raw RGB images (can be renamed via --raw_dir)
• data: JSON files containing DNG/JPEG metadata (can be renamed via --metadata_dir)
• dngs: Original DNG files (can be renamed via --dng_dir)
• srgb_images: Corresponding processed sRGB images (can be renamed via --srgb_dir)

After this, run the labeling tool from labeling_tool/labeling_tool.m and select the directory (i.e., --output_path) that contains the generated subdirectories.

📄 Citation

If you use this code or dataset in your research, please cite our paper:

@inproceedings{afifi2025time-aware-awb,
  title={Time-Aware Auto White Balance in Mobile Photography},
  author={Afifi, Mahmoud and Zhao, Luxi and Punnappurath, Abhijith and Abdelsalam, Mohamed Ashraf and Zhang, Ran and Brown, Michael S.},
  booktitle={The IEEE International Conference on Computer Vision (ICCV)},
  year={2025}
}

📬 Contact

For questions or inquiries, please contact:

• Mahmoud Afifi ([email protected], [email protected])
• Luxi Zhao ([email protected], [email protected])