MW-GAN

This repo is the official code for the following papers:

MW-GAN+ for Perceptual Quality Enhancement on Compressed Video. Jianyi Wang, Mai Xu (Corresponding), Xin Deng, Liquan Shen, Yuhang Song.

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

Multi-level Wavelet-based Generative Adversarial Network for Perceptual Quality Enhancement of Compressed Video. Jianyi Wang, Xin Deng, Mai Xu, Congyong Chen, Yuhang Song.

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

This repo is completely built based on BasicSR. You need to install following Install from a local clone. Quick 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:

MWGAN PSNR Model: Pretrained PSNR model to initialize generator for GAN training.
MWGAN GAN Model: The model for MW-GAN.

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

zzwei1 / MW-GAN