Official codes for High-Resolution Image Harmonization with a Simple Hybrid CNN-Transformer Network

Prerequisites

Linux
Python 3
CPU or NVIDIA GPU + CUDA CuDNN

Install basicsr

python setup_basicsr.py develop --no_cuda_ext

Train/Test

Download iHarmony4 dataset. We have partitioned a validation set from the original iHarmony4 dataset. Detailed partitioning text files can be found at Google Drive
Train our model:

CUDA_VISIBLE_DEVICES=0 python train.py --model SPANET --name experimental_name --dataset_root /***/iHarmony4/HAdobe5k/ --dataset_name HAdobe5k --batch_size 4 --init_port 55554 --local_rank 4 --crop_size 1024 --load_size 1024 --netG SPANET

Test our model:

CUDA_VISIBLE_DEVICES=0 python test.py --model SPANET --name experimental_name --dataset_root /***/iHarmony4/HAdobe5k/ --dataset_name HAdobe5k --batch_size 4 --init_port 55554 --local_rank 4 --crop_size 1024 --load_size 1024 --netG SPANET

Apply a pre-trained model

Download pre-trained models from

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 HAdobe5k  --image_size 1024

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, IntrinsicHarmony and DHT

guanguanboy / SHT-Net