search

fenglinglwb / EDT

On Efficient Transformer-Based Image Pre-training for Low-Level Vision
130 stars 9 forks source link

readme

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

  1. Clone the repository.

    git clone https://github.com/fenglinglwb/EDT.git

  2. Install the dependencies.

    • Python >= 3.7
    • PyTorch >= 1.4
    • Other packages 
      pip install -r requirements.txt

  3. 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

  4. 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}
}

Acknowledgement

We refer to Simple-SR and SwinIR for some details.