Skip to content
forked from IceClear/MW-GAN

Official PyTorch implements for Multi-level Wavelet-Based Generative Adversarial Network for Perceptual Quality Enhancement of Compressed Video (ECCV, 2020) and MW-GAN+ for Perceptual Quality Enhancement on Compressed Video (IEEE TCSVT).

License

Notifications You must be signed in to change notification settings

MC-2-Lab/MW-GAN

 
 

Repository files navigation

MW-GAN

This repo is the official code for the following papers:

Published on IEEE Transactions on Circuits and Systems for Video Technology in 2021. By MC2 Lab @ Beihang University.

Published on 16TH EUROPEAN CONFERENCE ON COMPUTER VISION in 2020. By MC2 Lab @ Beihang University.

Visual results on JCT-VC

Compressed video (QP=42) Ours
:-------------------------: :-------------------------:

Dependencies and Installation

pip install -r requirements.txt
BASICSR_EXT=True python setup.py develop

Dataset Preparation

Generally, we directly read cropped images from folders.

  • Run data_process.py to extract frames from videos.
  • This repo should also support LMDB format for faster IO speed as BasicSR. Not tested yet.

Get Started

The same as BasicSR, you can see here for details.

MWGAN+ Train:

  • MWGAN+ PSNR Model: CUDA_VISIBLE_DEVICES=0 python basicsr/train.py -opt options/train/MWGAN/train_MWGAN_PSNR.yml
  • MWGAN+ GAN Model: CUDA_VISIBLE_DEVICES=0 python basicsr/train.py -opt options/train/MWGAN/train_MWGAN_Percep.yml
  • Tradeoff Model: CUDA_VISIBLE_DEVICES=0 python basicsr/train.py -opt options/train/MWGAN/train_MWGAN_Tradeoff.yml

MWGAN Train:

  • MWGAN PSNR Model: CUDA_VISIBLE_DEVICES=0 python basicsr/train.py -opt options/train/MWGAN/train_MWGAN_ECCV_PSNR.yml
  • MWGAN GAN Model: CUDA_VISIBLE_DEVICES=0 python basicsr/train.py -opt options/train/MWGAN/train_MWGAN_ECCV_Percep.yml

Test:

  • Test example: CUDA_VISIBLE_DEVICES=0 python basicsr/test.py -opt options/test/MWGAN/test_MWGAN_Tradeoff.yml

Pre-train model

Here the models we provide are trained on QP37 in RGB space. You can also refer to Release.

MWGAN+ Model:

  • MWGAN+ PSNR Model: This is the model for MW-GAN+obj in the paper.
  • MWGAN+ GAN Model: This is the model for MW-GAN+ in the paper.
  • Tradeoff Model: For PD-tradeoff, instead of the ways introduced in our paper, we further developed an end-to-end model to achieve such a performance. Specifically, we first enhance the frames using the pre-trained PSNR-based model to remove compression artifacts, then using GAN to add high-frequency details. This two-stage enhancement is similar to the 'Two-stage Restoration' used in EDVR.

MWGAN Model:

Acknowledgement

This repo is built mainly based on BasicSR, and also borrow codes from pacnet and MWCNN_PyTorch. We thank a lot for their contributions to the community.

Citation

If you find our paper or code useful for your research, please cite:

@inproceedings{wang2020multi,
  title={Multi-level Wavelet-Based Generative Adversarial Network for Perceptual Quality Enhancement of Compressed Video},
  author={Wang, Jianyi and Deng, Xin and Xu, Mai and Chen, Congyong and Song, Yuhang},
  booktitle={European Conference on Computer Vision},
  pages={405--421},
  year={2020},
  organization={Springer}
}

@ARTICLE{wang2021mw,
  author={Wang, Jianyi and Xu, Mai and Deng, Xin and Shen, Liquan and Song, Yuhang},
  journal={IEEE Transactions on Circuits and Systems for Video Technology},
  title={MW-GAN+ for Perceptual Quality Enhancement on Compressed Video},
  year={2021},
  volume={},
  number={},
  pages={1-1},
  doi={10.1109/TCSVT.2021.3128275}
}

About

Official PyTorch implements for Multi-level Wavelet-Based Generative Adversarial Network for Perceptual Quality Enhancement of Compressed Video (ECCV, 2020) and MW-GAN+ for Perceptual Quality Enhancement on Compressed Video (IEEE TCSVT).

Resources

License

Stars

Watchers

Forks

Packages

No packages published

Languages

  • Python 89.1%
  • Cuda 6.5%
  • C++ 4.4%