Goal Conditioned Reinforcement Learning for Photo Finishing Tuning

Project Page | Paper | Video

This repo contains the code for Goal Conditioned Reinforcement Learning for Photo Finishing Tuning (NeurIPS'24)

---

Overview

RLPixTuner is an implementation of
"Goal Conditioned Reinforcement Learning for Photo Finishing Tuning"
Jiarui Wu, Yujin Wang, Lingen Li, Zhang Fan, Tianfan Xue
in Conference on Neural Information Processing Systems (NeurIPS), 2024

In this work, we propose an RL-based photo finishing tuning algorithm that efficiently tunes the parameters of a black-box photo finishing pipeline to match any tuning target. The RL-based solution (top row) takes only about 10 iterations to achieve a similar PSNR as the 500-iteration output of a zeroth-order algorithm (bottom row). Our method demonstrates fast convergence, high quality, and no need for a proxy.

Quick Start

Tested with python 3.9.

Environment

Clone our repo:

git clone https://github.com/OpenImagingLab/RLPixTuner.git
cd RLPixTuner

Install dependencies:

conda create --name <env> --file requirements.txt

Data

Prepare the input-target data pairs like:

├─ExpertC0001-0-Best-Input.tif
└─ExpertC0001-0-Best-Target.tif
...

Here is the an example (subset) of the FiveK-Random Evaluation Data we used in our paper.

After downloading the data, run:

python tools/process_data.py /path/to/your/dataset/directory

Model

The pretrained weight for photo finishing tuning comes with our repo:

envs/checkpoints/best_model.zip

Usage

This repo currently contains implementation for photo finishing tuning task in our paper.

Inference

Specify the eval data path and model weight in bash/run.sh. Then run:

bash bash/run.sh

Train

Specify the training data path in bash/train.sh. Then run:

bash bash/train.sh

Citation

If you use our work in your research, please use the following BibTeX entry.

@article{wu2024goal,
  title={Goal Conditioned Reinforcement Learning for Photo Finishing Tuning},
  author={Jiarui Wu and Yujin Wang and Lingen Li and Zhang Fan and Tianfan Xue},
  booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
  year={2024},
}

License

This repository is licensed under Attribution-NonCommercial 4.0 International

1. Attribution — Proper credit must be given, including a link to the license and an indication of any modifications made. This should be done in a reasonable manner, without implying endorsement by the licensor

2. NonCommercial — The Algorithm may NOT be used for commercial purposes. This includes, but is not limited to, the sale, licensing, or integration of the Software into commercial products or services.

For collaboration or inquiries, please contact us.

Acknowledgement

This code is based on the StableBaseline3.