CHROMA-MILL: Joint Low-Light Enhancement and White Balance under Colored Illumination

Official code for CTI-Retinexformer, the joint low-light enhancement and white-balance model of the WACV 2027 paper

Pol Riubrogent Comas, David Serrano-Lozano, Maria Pilligua, Michael S. Brown, Ramon Baldrich, Javier Vazquez-Corral. CHROMA-MILL: Joint Low-Light Enhancement and White Balance under Colored Illumination. WACV 2027.

Installation

conda create -n chromamill python=3.11 -y && conda activate chromamill
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128   # pick your CUDA
pip install -r requirements.txt

Quick test on the bundled samples

samples/nikon holds one test scene (Test15) at Levels 2, 5 and 10 under the four captured illuminants and two blends, plus its ground truth. samples/phone holds the same scene from the smartphone at Level 5.

python test.py --weights pretrained/CTI-Retinexformer_nikon.pth \
               --input samples/nikon/low --gt samples/nikon/high --output results/samples_nikon

results/samples_nikon/metrics.csv should match samples/nikon/expected_metrics.csv (the paper's per-image numbers) to within about 0.01. Small differences come from GPU vs. CPU arithmetic. The means over these 18 images are:

PSNR_L	PSNR_C	SSIM	LPIPS	MAE∠	ΔE76	ΔE00
20.873	20.463	0.763	0.142	9.205	15.231	11.145

The predictions are linear images, like the inputs. sRGB-encoded copies for viewing are written to preview/. predictions.csv lists the illuminant chromaticity and intensity the model reads from its own bottleneck (spatial means of Z_ch and Z_I).

Testing

On CHROMA-MILL. This scores every input against the White-illuminant Level-11 capture of its scene, as in the paper:

python test.py --weights pretrained/CTI-Retinexformer_nikon.pth \
               --input datasets/CHROMA-MILL/nikon/test/low \
               --gt datasets/CHROMA-MILL/nikon/test/high --output results/nikon

Use pretrained/CTI-Retinexformer_phone.pth with the smartphone split.

On your own images. For sRGB images (JPEG, PNG from a camera or phone), add --srgb. The inputs are linearized with the inverse sRGB curve before inference, and the outputs are re-encoded to sRGB. This is how Fig. 5 of the paper was produced. The models were trained only on CHROMA-MILL, so expect a domain gap.

python test.py --weights pretrained/CTI-Retinexformer_nikon.pth --input my_images/ --output results/mine --srgb

Dataset

Download CHROMA-MILL from [link to be added] and place it as

datasets/CHROMA-MILL/
  nikon/{train,val,test}/{low,high}/
  phone/{train,val,test}/{low,high}/

Every file is named <Scene>-F_<f>-SS_<shutter>-x_<XXXX>-y_<YYYY>-B_<B>.png. Here (XXXX, YYYY)/10000 is the CIE xy of the illuminant and B the light level. B = 254 is the reference (Level 11) and B ∈ {0, 80, 113, 139, 161, 180, 197, 213, 227, 241} are the ten low-light inputs (Levels 1–10). The final ground truth of every input is the White (x_3127-y_3290) B = 254 capture of its scene in high/.

python scripts/generate_combinations.py --src <captured>/nikon --dst datasets/CHROMA-MILL/nikon

Training

python train.py -opt options/train_cti_nikon.yml       # DSLR, from scratch
python train.py -opt options/train_cti_phone.yml       # smartphone, fine-tuned from the DSLR model

Outputs go to experiments/<name>/. The paper's models are the final checkpoints (weights/net_g_latest.pth) of these schedules: 3000 epochs at batch 32 (≈675k iterations, about 3 days on one GPU) for the DSLR, and 3000 epochs at batch 64 for the smartphone. best_psnr.pth (best validation PSNR) is also saved. Set logger.use_wandb: true to log to Weights & Biases.