On Efficient Transformer-Based Image Pre-training for Low-Level Vision
Wenbo Li*, Xin Lu*, Shengju Qian, Jiangbo Lu, Xiangyu Zhang, Jiaya Jia
[Paper]
News
We have made the testing code and well-trained models for SR, denoising and deraining available now. The training code will be released soon.
Usage
-
Clone the repository.
git clone https://github.com/fenglinglwb/EDT.git
-
Install the dependencies.
- Python >= 3.7
- PyTorch >= 1.4
- Other packages
pip install -r requirements.txt
-
Download pretrained models from One Drive and put them into folder 'pretrained'. Models are named by
task_model_data[__pretrain-task_pretrain-data]
where the optional part (in square brackets []) denotes the pre-training setting.
- Task
- Super-Resolution (SR) includes x2, x3, x4 scales.
- Denoising (DN) includes Gaussian noise levels 15, 25 and 50, i.e., g15, g25, g50.
- Deraining (DR) includees light and heavy streaks, i.e., ls and hs.
- Type
- Fine-tune: models are fine-tuned on target datasets with pre-training on ImageNet.
- Pre-train: models are trained on ImageNet.
- Scratch: models are trained on target datasets.
- Model
- EDT: T, S, B, L represent the tiny, small, base and large models.
- EDTSF: SF means the denoising or deraining model without downsampling and upsampling in the encoder and decoder.
- Datasets
- Pre-train: ImageNet.
- SR: Div2K, Flickr2K.
- Denoising: Div2K, Flickr2K, BSD500 and WED, short for D4.
- Deraining: RAIN100L, RAIN100H.
- Note. We only provide pre-trained and fine-tuned models for deraining since RAIN100 dataests only contain hundreds of low-resolution images (insufficient to train transformers).
Task Type Model SR Fine-tune SRx2_EDTT_Div2kFlickr2K__SRx2x3x4_ImageNet200K SRx3_EDTT_Div2kFlickr2K__SRx2x3x4_ImageNet200K SRx4_EDTT_Div2kFlickr2K__SRx2x3x4_ImageNet200K SRx2_EDTB_Div2kFlickr2K__SRx2x3x4_ImageNet200K SRx3_EDTB_Div2kFlickr2K__SRx2x3x4_ImageNet200K SRx4_EDTB_Div2kFlickr2K__SRx2x3x4_ImageNet200K SRx2_EDTT_Div2kFlickr2K__SRx2_ImageNet200K SRx2_EDTS_Div2kFlickr2K__SRx2_ImageNet200K SRx2_EDTB_Div2kFlickr2K__SRx2_ImageNet200K SRx3_EDTB_Div2kFlickr2K__SRx3_ImageNet200K SRx4_EDTB_Div2kFlickr2K__SRx4_ImageNet200K SRx2_EDTL_Div2kFlickr2K__SRx2_ImageNet200K SRx2_EDTB_Div2kFlickr2K__SRx2_ImageNet50K SRx2_EDTB_Div2kFlickr2K__SRx2_ImageNet100K SRx2_EDTB_Div2kFlickr2K__SRx2_ImageNet400K SRx2_EDTB_Div2kFlickr2K__SRx2_ImageNetFull Pre-train SRx2x3x4_EDTT_ImageNet200K SRx2x3x4_EDTB_ImageNet200K SRx2x3DNg15_EDTB_ImageNet200K SRx2_EDTT_ImageNet200K SRx2_EDTS_ImageNet200K SRx2_EDTB_ImageNet200K SRx3_EDTB_ImageNet200K SRx4_EDTB_ImageNet200K SRx2_EDTL_ImageNet200K SRx2_EDTB_ImageNet50K SRx2_EDTB_ImageNet100K SRx2_EDTB_ImageNet400K SRx2_EDTB_ImageNetFull Scratch SRx2_EDTT_Div2kFlickr2K SRx3_EDTT_Div2kFlickr2K SRx4_EDTT_Div2kFlickr2K SRx2_EDTS_Div2kFlickr2K SRx2_EDTB_Div2kFlickr2K SRx3_EDTB_Div2kFlickr2K SRx4_EDTB_Div2kFlickr2K SRx2_EDTL_Div2kFlickr2K Denoise Fine-tune DNg15_EDTB_D4__DNg15g25g50_ImageNet200K DNg25_EDTB_D4__DNg15g25g50_ImageNet200K DNg50_EDTB_D4__DNg15g25g50_ImageNet200K DNg15_EDTB_D4__DNg15_ImageNet200K DNg25_EDTB_D4__DNg25_ImageNet200K DNg50_EDTB_D4__DNg50_ImageNet200K Pre-train DNg15g25g50_EDTB_ImageNet200K DNg15_EDTB_ImageNet200K DNg25_EDTB_ImageNet200K DNg50_EDTB_ImageNet200K Scratch DNg15_EDTB_D4 DNg25_EDTB_D4 DNg50_EDTB_D4 DNg15_EDTBSF_D4 DNg25_EDTBSF_D4 DNg50_EDTBSF_D4 Derain Fine-tune DRls_EDTB_RAIN100L__DRlshs_ImageNet200K DRhs_EDTB_RAIN100H__DRlshs_ImageNet200K DRls_EDTB_RAIN100L__DRls_ImageNet200K DRhs_EDTB_RAIN100H__DRhs_ImageNet200K Pre-train DRlshs_EDTB_ImageNet200K DRls_EDTB_ImageNet200K DRhs_EDTB_ImageNet200K - Task
-
Quick test.
The model and config files are in one-to-one correpondence with the same name. Please refer to the naming rule in the model zoo above.
-
SR and deraining.
Read low-quality data directly from a specified folder as
python test_sample.py --config config_path --model model_path --input input_folder [--output output_folder --gt gt_folder]
where '--output' and '--gt' are optional. If assigned, the predictions will be stored and PSNR/SSIM results will be reported.
For example,
python test_sample.py --config configs/SRx2_EDTT_Div2kFlickr2K.py --model pretrained/SRx2_EDTT_Div2kFlickr2K.pth --input test_sets/SR/Set5/LR/x2 --gt test_sets/SR/Set5/HR/x2
-
Denoising.
The low-quality data is obtained by adding noise to the gt as
python test_sample.py --config config_path --model model_path --gt gt_folder --noise_level XX [--output output_folder --sf]
where '--sf' indicates whether there is upsampling and downsampling. If not assigned, EDT model will be built.
For example,
python test_sample.py --config configs/DNg15_EDTB_D4.py --model pretrained/DNg15_EDTB_D4.pth --gt test_sets/Denoise/McMaster --noise_level 15
-
Note.
The pre-training may contain multiple tasks. If you want to test multi-task models, please only build one branch and load corresponding weights during model building phase. We have provided an example for testing x2 SR based on model 'SRx2x3x4_EDTB_ImageNet200K' in the comment of 'test_sample.py'.
-
Citation
@article{li2021efficient,
title={On Efficient Transformer and Image Pre-training for Low-level Vision},
author={Li, Wenbo and Lu, Xin and Qian, Shengju and Lu, Jiangbo and Zhang, Xiangyu and Jia, Jiaya},
journal={arXiv preprint arXiv:2112.10175},
year={2021}
}