A Pytorch Implementation of a continuously rate adjustable learned image compression framework, Gained Variational Autoencoder(GainedVAE).
Note that It Is Not An Official Implementation Code.
More details can be found in the following paper:
Asymmetric Gained Deep Image Compression With Continuous Rate Adaptation.
Huawei Technologies, CVPR 2021
Ze Cui, Jing Wang, Shangyin Gao, Tiansheng Guo, Yihui Feng, Bo Bai
The differences with the original paper are not limited to the following:
- The number of model channels may be different.
- The pre-defined lambda set may be different.
- Asymmetric Gaussian entropy model is not implemented.
-------------------------------2023.4.6 Update--------------------------------
- eval script is uploaded.
- A spatial-channel feature modulation framework with Gain Units is Added in compressai.models.gain . You could find more details in [1]1.
- Python == 3.7.10
- Pytorch == 1.7.1
- CompressAI
I use a part of the OpenImages Dataset to train the models (train06, train07, train08, about 54w images). You can download from hereDownload OpenImages. Maybe train08 (14w images) is enough.
python3 trainGain.py -d /path/to/your/image/dataset/ --epochs 200 -lr 1e-4 --batch-size 16 --model-save /path/to/your/model/save/dir --cuda
Currently only two eval modes is supported, one is 'gain' and the other is 'scgain'.
python3 eval_gain.py -d /path/to/your/image/dataset/ --checkpoint /path/to/your/model.pth --logpath /path/to/save/result/log/ --cuda --mode (gain/scgain)
I try to train the Gained Mean-Scale Hyperprior model/Gained Scale Hyperprior model. See details in ./results I retrained the single rate baseline but can not achieve the official performance. Results from Google tensorflow/compression library is very strong probably because of their large and diverse training data set and long training time.
You can download the checkpoint trained by me from Pretrained Model.
The framework is based on CompressAI, I add the model in compressai.models.gain, compressai.models.gain_utils.
And trainGain/trainGain.py is modified with reference to compressai_examples/train.py.
[1] "Variable-Rate Deep Image Compression through Spatially-Adaptive Feature Transform" (ICCV 2021). code
[2] "Variable Bitrate Image Compression with Quality Scaling Factors" (ICASSP 2020). code
[3] "Variable Rate Deep Image Compression with Modulated Autoencoders" (IEEE SPL 2020) code
[4] "Slimmable Compressive Autoencoders for Practical Neural Image Compression" (CVPR 2021) code
This work can be regarded as the feature modulation version of GainVAE and can be easily implemented based on this repo.
Feel free to contact me if there is any question about the code or to discuss any problems with image and video compression. (mxh_wine@qq.com)