Core codes for DUAL-HEAD FUSION NETWORK FOR IMAGE ENHANCEMENT
Framework of our proposed method which consists of
- a dual-head feature extraction module
- a global color rendering module
- a context-aware retouching module
- a spatial attention based fusion
All modules are jointly learned from the annotated data in an end-to-end manner.
pip install -r requirements.txt
- MIT-Adobe FiveK Dataset
- We use the setting of Image-Adaptive-3DLUT in our experiments, please refer to their page for details and data link.
- PPR10K
Prepare the dataset in the following format and you could use the provided FiveK Dataset class.
- <data_root>
- input_train
- input_test
- target_train
- target_test
Or you need to implement your own Class for your customed data format / directory arrangement.
The default settings of the most hyper-parameters are written in the parameter.py file. To get started as soon as possible (with the FiveK dataset), only the 'data_root' needs to be modified before training.
python train.py --data_root <path>
By default, the images, models, and logs generated during training are saved in save_root/dataset/name.
To evaluate your own trained model of a specific epoch, specify the epoch and keep the other parameters the same as training.
For example,
python evaluate.py --model *** --epoch 397
-
For video ./video