Foreground-aware Semantic Representations for Image Harmonization

This repository contains the official PyTorch implementation of the following paper:

Foreground-aware Semantic Representations for Image Harmonization
Konstantin Sofiiuk, Polina Popenova, Anton Konushin
Samsung AI Center Moscow
https://arxiv.org/abs/2006.00809

Abstract: Image harmonization is an important step in photo editing to achieve visual consistency in composite images by adjusting the appearances of foreground to make it compatible with background. Previous approaches to harmonize composites are based on training of encoder-decoder networks from scratch, which makes it challenging for a neural network to learn a high-level representation of objects. We propose a novel architecture to utilize the space of high-level features learned by a pre-trained classification network. We create our models as a combination of existing encoder-decoder architectures and a pre-trained foreground-aware deep high-resolution network. We extensively evaluate the proposed method on existing image harmonization benchmark and set up a new state-of-the-art in terms of MSE and PSNR metrics.

Setting up an environment

This framework is built using Python 3.6 and relies on the PyTorch 1.4.0+. The following command installs all necessary packages:

pip3 install -r requirements.txt

You can also use our Dockerfile to build a container with configured environment.

If you want to run training or testing, you must configure the paths to the datasets in config.yml.

Datasets

We train and evaluate all our models on the iHarmony4 dataset. It contains 65742 training and 7404 test objects. Each object is a triple consisting of real image, composite image and foreground mask.

Before training we resize HAdobe5k subdataset so that each side is smaller than 1024. The resizing script is provided in resize_hdataset.ipynb.

Don't forget to change the paths to the datasets in config.yml after downloading and unpacking.

Training

We provide the scripts for training our models on images of size 256 and 512. For each experiment, a separate folder is created in the

./harmonization_exps

EXPS_PATH

Start training with the following commands: ```.bash python3 train.py --gpus=0 --workers=4 --exp-name=first-try

iDIH: fully convolutional Encoder-Decoder with output image blending and foreground-normalized MSE loss

python3 train.py models/fixed256/improved_dih.py --gpus=0 --workers=4 --exp-name=first-try

HRNet18s-V2p + iDIH: feature pyramid of 4 HRNet18-small-V2 outputs is concatenated to 4 outputs of the iDIH encoder

python3 train.py models/fixed256/hrnet18_idih.py --gpus=0 --workers=4 --exp-name=first-try

HRNet18-V2 + iDIH: single output of HRNet18-V2 is concatenated to single output of the iDIH encoder

python3 train.py models/fixed256/hrnet18_idih.py --gpus=0 --workers=4 --exp-name=first-try

iDIH trained on 512x512

python3 train.py models/crop512/improved_dih.py --gpus=0 --workers=4 --exp-name=first-try ``

To see all training parameters, run

We used pre-trained HRNetV2 models from the official repository. To train one of our models with HRNet backbone, download HRNet weights and specify their path in config.yml (see

IMAGENET_PRETRAINED_MODELS

Evaluation

We provide scripts to both evaluate and get predictions from any model. To do that, we specify all models configs in mconfigs. To evaluate a model different from the provided, a new config entry should be added.

You can specify the checkpoints path in config.yml (see

MODELS_PATH

Evaluate model

To get metrics table on the iHarmony4 test set run the following command: ```.bash python3 scripts/evaluate_model.py --resize-strategy Fixed256

iDIH

python3 scripts/evaluatemodel.py improveddih256 /hdd0/harmonizationexps/fixed256/improveddih/checkpoints/lastcheckpoint.pth --resize-strategy Fixed256 ``

To see all evaluation parameters run

Get model predictions

To get predictions on a set of images, run the following command: ```.bash python3 scripts/predictfordir.py --images --masks --resize 256

iDIH

python3 scripts/evaluatemodel.py improveddih256 /hdd0/harmonizationexps/fixed256/improveddih/checkpoints/lastcheckpoint.pth \ --images /hdd0/datasets/ImageHarmonization/test/compositeimages --masks /hdd0/datasets/ImageHarmonization/test/masks \ --resize 256 ``

To see all evaluation parameters run

Jupyter notebook

For interactive models testing with samples visualization see evalandvisharmonizationmodel.ipynb.

Results

We provide metrics and pre-trained weights for several models trained on images of size 256x256 augmented with horizontal flip and random resized crop. Metric values may differ slightly from the ones in the paper since all the models were retrained from scratch with the new codebase.

Pre-trained models with corresponding names of model configs (see Evaluation): | Model | Download Link | Name in mconfigs | |:--------------------------:|:----------------------------------------------:|:-------------------------------------:| | iDIH256 | idih256.pth |improveddih256 | | iSSAM256 | issam256.pth |improvedssam256 | | DeepLab-ResNet34 + iDIH256 | deeplab_idih256.pth |deeplabr34idih256 | | HRNet18s + iDIH256 | hrnet18s_idih256.pth |hrnet18sidih256 | | HRNet18 + iDIH256 | hrnet18_idih256.pth |hrnet18idih256 | | HRNet18 pyramid + iDIH256 | hrnet18v2pidih256.pth |hrnet18v2pidih256 | | HRNet32 + iDIH256 | hrnet32_idih256.pth |hrnet32_idih256 |

Evaluation metrics:

License

The code is released under the MPL 2.0 License. MPL is a copyleft license that is easy to comply with. You must make the source code for any of your changes available under MPL, but you can combine the MPL software with proprietary code, as long as you keep the MPL code in separate files.

Citation

If you find this work is useful for your research, please cite our paper:

@article{sofiiuk2020harmonization,
  title={Foreground-aware Semantic Representations for Image Harmonization},
  author={Konstantin Sofiiuk, Polina Popenova, Anton Konushin},
  journal={arXiv preprint arXiv:2006.00809},
  year={2020}
}