[CVPR 2026] CanonCGT: Reference-Based Color Grading via Canonical Pivot Representation.

Jinwon Ko, Keunsoo Ko, and Chang-Su Kim.

Official code for "CanonCGT: Reference-Based Color Grading via Canonical Pivot Representation" in CVPR 2026. [arXiv] [paper] [video]

📝 Introduction

We present CanonCGT, a reference-based color grading framework based on a canonical pivot representation. Our key idea is to first map the input image into a style-neutral canonical domain and then apply the reference-driven grading style from this canonical representation.

We also introduce DP-CGT, a dual-phase training strategy that combines supervised preset learning and self-supervised refinement for robust generalization to diverse reference images.

⚙️ Preparation

1. Installation

Create conda environment:

$ conda create -n CanonCGT python=3.9 anaconda
$ conda activate CanonCGT
$ conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
$ pip install opencv-python-headless==4.10.0.82
$ pip install pyyaml scikit-learn lpips

2. Pretrained models

Pretrained models will be available in:

root/CanonCGT/pretrained/

They can also be downloaded from here.

🚀 Demo

You can run a demo with pretrained models to perform reference-based color grading on your own images.

$ cd root/CanonCGT/
$ python demo.py \
    --gpu 0 \
    --pretrained_path ./pretrained/SSL.pth \
    --inp_path SET_YOUR_INPUT_PATH \
    --ref_path SET_YOUR_REFERENCE_PATH \
    --out_path SET_YOUR_OUTPUT_PATH

The color-graded result will be saved to SET_YOUR_OUTPUT_PATH

📂 Dataset

We use both supervised paired data and unsupervised image collections for training CanonCGT.

1. Supervised Paired Dataset

For supervised preset learning, we construct a paired dataset from the MIT-Adobe FiveK dataset. Specifically, we use the Expert C version as the canonical target and generate preset-transformed images using Lightroom presets.

Due to license restrictions, we do not redistribute the Lightroom presets or the generated paired dataset. Instead, we provide the source links to the presets used in our experiments, so that users can download them directly from the original providers and follow their respective license terms.

• MIT-Adobe FiveK Dataset
• Lightroom preset source links

2. Unsupervised Dataset

For self-supervised refinement, we use diverse unpaired image datasets from multiple domains. Please download each dataset from its official source:

• Flickr2K
• DIV2K
• LSDIR
• PPR10K
• Food-101
• Google Landmarks Dataset v2

🏋️ Train

CanonCGT is trained using DP-CGT, a dual-phase training strategy consisting of supervised preset learning and self-supervised refinement.

Phase 1: Supervised Preset Learning

In Phase 1, CanonCGT is trained using the supervised paired dataset constructed from MIT-Adobe FiveK and Lightroom presets. This phase learns the canonical pivot representation and the reference-based grading process from preset-based paired data.

Phase 1-A: Grade Extractor Training

First, train the grade extractor to encode tonal characteristics from preset-transformed images.

$ cd root/CanonCGT/
$ python main.py \
    --gpu 0 \
    --yaml Stage1_style_encoder \
    --run_mode train

After training the grade extractor, save the style centroids:

$ python main.py \
    --gpu 0 \
    --yaml Stage1_style_encoder \
    --run_mode test \
    --load

Phase 1-B: Canonicalizer and Grader Training

Next, train the canonicalizer and the grader using the supervised paired dataset.

$ python main.py \
    --gpu 0 \
    --yaml Stage1_canonicalizer \
    --run_mode train

$ python main.py \
    --gpu 0 \
    --yaml Stage1_styler \
    --run_mode train

Phase 1-C: End-to-End Fine-tuning

Then, fine-tune the full CanonCGT framework in an end-to-end manner.

$ python main.py \
    --gpu 0 \
    --yaml Stage2_end_to_end_finetuning \
    --run_mode train

Phase 2: Self-Supervised Refinement

In Phase 2, CanonCGT is further refined using unpaired image datasets. This phase improves the generalization ability of CanonCGT to diverse real-world reference images beyond the preset-based supervised training data.

$ python main.py \
    --gpu 0 \
    --yaml Stage3_SSL_training_Flickr2K_PPR10K_LSDIR \
    --run_mode train

Alternatively, you can run the full training pipeline with:

$ bash run.sh

🧪 Evaluation

You can evaluate CanonCGT using a pretrained or trained model with:

$ python main.py \
    --gpu 0 \
    --yaml Stage3_SSL_training_Flickr2K_PPR10K_LSDIR \
    --run_mode eval \
    --load

The evaluation results will be saved to the output directory specified in the corresponding configuration file.

🎨 Results

The figure below shows our color grading results. For each pair, the left image shows the input and the right image shows the color-graded output using the inset reference image. CanonCGT produces photorealistic color grading that matches the tonal mood, lighting, and color temperature of the reference while preserving color harmony and scene structure.

The figures below provide qualitative comparisons with existing methods. The green and cyan boxes denote the reference and ground-truth images, respectively. The input is generated by applying a random color perturbation to the ground-truth image. Results from competing methods and CanonCGT are shown with their corresponding error maps below.

License

This repository, including the source code and pretrained weights, is released under the Apache License 2.0.

Copyright 2026 Jinwon Ko.

Citation

@inproceedings{ko2026canoncgt,
  title={CanonCGT: Reference-Based Color Grading via Canonical Pivot Representation},
  author={Ko, Jinwon and Ko, Keunsoo and Kim, Chang-Su},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  year={2026}
}