Towards Ultra-Resolution Neural Style Transfer via Thumbnail Instance Normalization (AAAI'2022)

Official PyTorch implementation for our URST (Ultra-Resolution Style Transfer) framework.

URST is a versatile framework for ultra-high resolution style transfer under limited GPU memory resources, which can be easily plugged in most existing neural style transfer methods.

With the growth of the input resolution, the memory cost of our URST hardly increases. Theoretically, it supports style transfer of arbitrary resolution images.

One ultra-high resolution stylized result of 12000 x 8000 pixels (i.e., 96 megapixels).

This repository is developed based on six representative style transfer methods, which are Johnson et al., MSG-Net, AdaIN, WCT, LinearWCT, and Wang et al. (Collaborative Distillation).

For details see Towards Ultra-Resolution Neural Style Transfer via Thumbnail Instance Normalization.

If you use this code for a paper please cite:

@inproceedings{chen2022towards,
  title={Towards Ultra-Resolution Neural Style Transfer via Thumbnail Instance Normalization},
  author={Chen, Zhe and Wang, Wenhai and Xie, Enze and Lu, Tong and Luo, Ping},
  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
  year={2022}
}

Environment

python3.6, pillow, tqdm, torchfile, pytorch1.1+ (for inference)

pip install pillow
pip install tqdm
pip install torchfile
conda install pytorch==1.1.0 torchvision==0.3.0 -c pytorch

tensorboardX (for training)
```
pip install tensorboardX
```

Then, clone the repository locally:

git clone https://github.com/czczup/URST.git

Test (Ultra-high Resolution Style Transfer)

Step 1: Prepare images

Content images and style images are placed in examples/.
Since the ultra-high resolution images are quite large, we not place them in this repository. Please download them from this google drive.
All content images used in this repository are collected from pexels.com.

Step 2: Prepare models

Download models from this google drive. Unzip and merge them into this repository.

Step 3: Stylization

First, choose a specific style transfer method and enter the directory.

Then, please run the corresponding script. The stylized results will be saved in output/.

For Johnson et al., we use the PyTorch implementation Fast-Neural-Style-Transfer.

cd Johnson2016Perceptual/
CUDA_VISIBLE_DEVICES=<gpu_id> python test.py --content <content_path> --model <model_path> --URST

For MSG-Net, we use the official PyTorch implementation PyTorch-Multi-Style-Transfer.

cd Zhang2017MultiStyle/
CUDA_VISIBLE_DEVICES=<gpu_id> python test.py --content <content_path> --style <style_path> --URST

For AdaIN, we use the PyTorch implementation pytorch-AdaIN.

cd Huang2017AdaIN/
CUDA_VISIBLE_DEVICES=<gpu_id> python test.py --content <content_path> --style <style_path> --URST

For WCT, we use the PyTorch implementation PytorchWCT.

cd Li2017Universal/
CUDA_VISIBLE_DEVICES=<gpu_id> python test.py --content <content_path> --style <style_path> --URST

For LinearWCT, we use the official PyTorch implementation LinearStyleTransfer.

cd Li2018Learning/
CUDA_VISIBLE_DEVICES=<gpu_id> python test.py --content <content_path> --style <style_path> --URST

For Wang et al. (Collaborative Distillation), we use the official PyTorch implementation Collaborative-Distillation.

cd Wang2020Collaborative/PytorchWCT/
CUDA_VISIBLE_DEVICES=<gpu_id> python test.py --content <content_path> --style <style_path> --URST

For Multimodal Transfer, we use the PyTorch implementation multimodal_style_transfer

cd Wang2017Multimodal/
CUDA_VISIBLE_DEVICES=<gpu_id> python test.py --content <content_path> --model <model_name> --URST

Optional options:

--patch_size: The maximum size of each patch. The default setting is 1000.
--style_size: The size of the style image. The default setting is 1024.
--thumb_size: The size of the thumbnail image. The default setting is 1024.
--URST: Use our URST framework to process ultra-high resolution images.

Train (Enlarge the Stroke Size)

Step 1: Prepare datasets

Download the MS-COCO 2014 dataset and WikiArt dataset.

MS-COCO

wget http://msvocds.blob.core.windows.net/coco2014/train2014.zip

WikiArt
- Either manually download from kaggle.
- Or install kaggle-cli and download by running:
```
kg download -u <username> -p <password> -c painter-by-numbers -f train.zip
```

Step 2: Prepare models

As same as the Step 2 in the test phase.

Step 3: Train the decoder with our stroke perceptual loss

For AdaIN:

cd Huang2017AdaIN/
CUDA_VISIBLE_DEVICES=<gpu_id> python trainv2.py --content_dir <coco_path> --style_dir <wikiart_path>

For LinearWCT:

cd Li2018Learning/
CUDA_VISIBLE_DEVICES=<gpu_id> python trainv2.py --contentPath <coco_path> --stylePath <wikiart_path>

License

This repository is released under the Apache 2.0 license as found in the LICENSE file.