Home

Awesome

PWC PWC

Pre-Trained Image Processing Transformer (IPT)

By Hanting Chen, Yunhe Wang, Tianyu Guo, Chang Xu, Yiping Deng, Zhenhua Liu, Siwei Ma, Chunjing Xu, Chao Xu, Wen Gao. [arXiv]

We study the low-level computer vision task (such as denoising, super-resolution and deraining) and develop a new pre-trained model, namely, image processing transformer (IPT). We present to utilize the well-known ImageNet benchmark for generating a large amount of corrupted image pairs. The IPT model is trained on these images with multi-heads and multi-tails. The pre-trained model can therefore efficiently employed on desired task after fine-tuning. With only one pre-trained model, IPT outperforms the current state-of-the-art methods on various low-level benchmarks.

<p align="center"> <img src="images/intro.png" width="900"> </p>

MindSpore Code

Requirements

Dataset

The benchmark datasets can be downloaded as follows:

For super-resolution:

Set5, Set14, B100, Urban100.

Please refer to EDSR-PyTorch

For denoising:

CBSD68, Urban100.

For deraining:

Rain100.

Please refer to PReNet

The result images are converted into YCbCr color space. The PSNR is evaluated on the Y channel only.

Script Description

This is the inference script of IPT, you can following steps to finish the test of image processing tasks, like SR, denoise and derain, via the corresponding pretrained models.

Script Parameter

For details about hyperparameters, see option.py.

Evaluation

Pretrained models

The pretrained models are available in google drive

Evaluation Process

Inference example: For SR x2,x3,x4:

python main.py --dir_data $DATA_PATH --pretrain $MODEL_PATH --data_test Set5+Set14+B100+Urban100 --scale $SCALE --test_only

Note: path of images should be like: $MODEL_PATH/benchmark/Set5/HR/XXX.png and $MODEL_PATH/benchmark/Set5/LR_bicubic/XXX.png

For Denoise 30,50:

python main.py --dir_data $DATA_PATH --pretrain $MODEL_PATH --data_test CBSD68+Urban100 --scale 1 --denoise --sigma $NOISY_LEVEL --test_only

Note: path of images should be like: $MODEL_PATH/benchmark/CBSD68/XXX.png

For derain:

python main.py --dir_data $DATA_PATH --pretrain $MODEL_PATH --scale 1 --derain --test_only

Note: path of images should be like: $MODEL_PATH/Rain100L/rain-XXX.png and $MODEL_PATH/Rain100L/norain-XXX.png

Results

<center>
MethodScaleSet5Set14B100Urban100
VDSRX237.5333.0531.9030.77
EDSRX238.1133.9232.3232.93
RCANX238.2734.1232.4133.34
RDNX238.2434.0132.3432.89
OISR-RK3X238.2133.9432.3633.03
RNANX238.1733.8732.3232.73
SANX2<u>38.31</u>34.07<u>32.42</u>33.1
HANX238.27<u>34.16</u>32.41<u>33.35</u>
IGNNX238.2434.0732.4133.23
IPT (ours)X238.3734.4332.4833.76
MethodScaleSet5Set14B100Urban100
VDSRX333.6729.7828.8327.14
EDSRX334.6530.5229.2528.80
RCANX334.7430.6529.3229.09
RDNX334.7130.5729.2628.80
OISR-RK3X334.7230.5729.2928.95
RNANX334.6630.5229.2628.75
SANX3<u>34.75</u>30.59<u>29.33</u>28.93
HANX3<u>34.75</u><u>30.67</u>29.32<u>29.10</u>
IGNNX334.7230.6629.3129.03
IPT (ours)X334.8130.8529.3829.49
MethodScaleSet5Set14B100Urban100
VDSRX431.3528.0227.2925.18
EDSRX432.4628.8027.7126.64
RCANX4<u>32.63</u>28.8727.7726.82
SANX432.64<u>28.92</u>27.7826.79
RDNX432.4728.8127.7226.61
OISR-RK3X432.5328.8627.7526.79
RNANX432.4928.8327.7226.61
HANX432.6428.90<u>27.80</u><u>26.85</u>
IGNNX432.5728.8527.7726.84
IPT (ours)X432.6429.0127.8227.26
</center> <p align="center"> <img src="images/sr_result.PNG" width="950"> </p> <p align="center"> <img src="images/dn_result.PNG" width="950"> </p> <p align="center"> <img src="images/dr_result.png" width="950"> </p>

Citation

@misc{chen2020pre,
      title={Pre-Trained Image Processing Transformer}, 
      author={Chen, Hanting and Wang, Yunhe and Guo, Tianyu and Xu, Chang and Deng, Yiping and Liu, Zhenhua and Ma, Siwei and Xu, Chunjing and Xu, Chao and Gao, Wen},
      year={2021},
      eprint={2012.00364},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

Acknowledgement