Awesome
<base target="_blank"/>Transformer for Image Harmonization and Beyond [Paper]<br>
Zonghui Guo, Zhaorui Gu, Bing Zheng, Junyu Dong, Haiyong Zheng<br> IEEE Transactions on Pattern Analysis and Machine Intelligence<br>
Here we provide the PyTorch implementation and pre-trained model of our latest version, if you require the code of our previous ICCV version ("Image Harmonization With Transformer"), please click the released version.
Prerequisites
- Linux
- Python 3
- CPU or NVIDIA GPU + CUDA CuDNN
Train/Test
-
Download iHarmony4 dataset.
-
Train our HT+ model (FC-TRE-DeCNN):
CUDA_VISIBLE_DEVICES=0 python train.py --model ht --tr_r_enc_head x --tr_r_enc_layers x --name experiment_name --dataset_root <dataset_dir> --dataset_name IHD --batch_size xx --init_port xxxx
- Test our HT+ model (FC-TRE-DeCNN):
CUDA_VISIBLE_DEVICES=0 python test.py --model ht --tr_r_enc_head x --tr_r_enc_layers x --name experiment_name --dataset_root <dataset_dir> --dataset_name IHD --batch_size xx --init_port xxxx
- Train our DHT+ model:
CUDA_VISIBLE_DEVICES=0 python train.py --model dht --light_use_mask --tr_r_enc_head x --tr_r_enc_layers x --tr_i_dec_head x --tr_i_dec_layers x --tr_l_dec_head x --tr_l_dec_layers x --name DHT_experiment_name --dataset_root <dataset_dir> --dataset_name IHD --batch_size xx --init_port xxxx
- Test our DHT+ model:
CUDA_VISIBLE_DEVICES=0 python test.py --model dht --light_use_mask --tr_r_enc_head x --tr_r_enc_layers x --tr_i_dec_head x --tr_i_dec_layers x --tr_l_dec_head x --tr_l_dec_layers x --name DHT_experiment_name --dataset_root <dataset_dir> --dataset_name IHD --batch_size xx --init_port xxxx
Apply a pre-trained model
- Download pre-trained models from Google Drive or BaiduCloud (access code: vmrg), and put
latest_net_G.pth
in the directorycheckpoints/HT_2H9L_allihd
orcheckpoints/DHT_2H9L_allihd
. Run:
# Our HT model
CUDA_VISIBLE_DEVICES=0 python test.py --model ht --tr_r_enc_head 2 --tr_r_enc_layers 9 --name HT_2H9L_allihd --dataset_root <dataset_dir> --dataset_name IHD --batch_size xx --init_port xxxx
# Our CNN-DHT model
CUDA_VISIBLE_DEVICES=0 python test.py --model dht --light_use_mask --tr_r_enc_head 2 --tr_r_enc_layers 9 --tr_i_dec_head 2 --tr_i_dec_layers 9 --tr_l_dec_head 2 --tr_l_dec_layers 9 --name DHT_2H9L_allihd --dataset_root <dataset_dir> --dataset_name IHD --batch_size xx --init_port xxxx
Evaluation
We provide the code in ih_evaluation.py
. Run:
CUDA_VISIBLE_DEVICES=0 python evaluation/ih_evaluation.py --dataroot <dataset_dir> --result_root results/experiment_name/test_latest/images/ --evaluation_type our --dataset_name ALL
Real composite image harmonnization
More compared results can be found at Google Drive or BaduCloud (access code: n37b).
Bibtex
If you use this code for your research, please cite our papers.
@article{guo2022transformer,
title={Transformer for Image Harmonization and Beyond},
author={Guo, Zonghui and Gu, Zhaorui and Zheng, Bing and Dong, Junyu and Zheng, Haiyong},
journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
year={2022},
publisher={IEEE}
}
@InProceedings{Guo_2021_ICCV,
author = {Guo, Zonghui and Guo, Dongsheng and Zheng, Haiyong and Gu, Zhaorui and Zheng, Bing and Dong, Junyu},
title = {Image Harmonization With Transformer},
booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
month = {October},
year = {2021},
pages = {14870-14879}
}
Acknowledgement
For some of the data modules and model functions used in this source code, we need to acknowledge the repositories of DoveNet, CycleGAN, SpiralNet and IntrinsicHarmony.