Home

Awesome

Neural Image Compression: Generalization, Robustness, and Spectral Biases

Official code implementation for NeurIPS 2023 paper "Neural Image Compression: Generalization, Robustness, and Spectral Biases".

Environment

Details about the environment are specified in environment.yml.

Data (CLIC, Kodak, and -C variants)

The 2020 Challenge on Learned Image Compression (CLIC) dataset can be downloaded from these links:

https://data.vision.ee.ethz.ch/cvl/clic/mobile_train_2020.zip https://data.vision.ee.ethz.ch/cvl/clic/professional_train_2020.zip https://data.vision.ee.ethz.ch/cvl/clic/mobile_valid_2020.zip https://data.vision.ee.ethz.ch/cvl/clic/professional_valid_2020.zip https://data.vision.ee.ethz.ch/cvl/clic/test/CLIC2020Mobile_test.zip https://data.vision.ee.ethz.ch/cvl/clic/test/CLIC2020Professional_test.zip

The Kodak dataset can be found here: https://r0k.us/graphics/kodak/.

The -C datasets were generated using code from: https://github.com/hendrycks/robustness/tree/master.

Training NIC models

train.py contains code to train the following neural image compression (NIC) models.

By default models are trained to optimize PSNR, but they can be trained to optimize MS-SSIM using the --distortion-metric flag. The training datasets and loaders are compatible with the CLIC dataset. See arguments in args/base_args.py and args/train_args.py for more details.

Example commands

Train a fixed-rate scale-hyperprior model:

python train.py --lmbda 0.05 --run-dir runs/sh_fr_dense/lmbda-0.05 --epochs 5000

Train an ELIC model:

python train.py --model elic  --N 192 --M 320 --lmbda 0.05 --run-dir runs/elic/lmbda-0.05 --epochs 5000 

Train a variable-rate scale-hyperprior model:

python train.py --run-dir runs/sh_vr_dense --variable-rate --lambda-range 0.001 0.25 --epochs 10000

Train a 95% sparse scale-hyperprior model:

python train.py --lmbda 0.05 --run-dir runs/sh_fr_pruned/sparsity-0.95/lmbda-0.05 --prune-algorithm gmp --final-sparsity 0.95

Testing models

test.py contains code to test image compression codecs.

Example commands

Test pre-trained SH NIC model on CLIC dataset:

python test.py --run-dir runs/sh_fr_dense/lmbda-0.05 --resume runs/sh_fr_dense/lmbda-0.05/checkpoints/epoch_4999.state --test-data clic --corruption 0 --severity 0 --psd

Test JPEG2000 on shot noise corruption of Kodak-C dataset:

python test.py --run-dir runs/jpeg2k/quality-20 --classic-codec jpeg2k -q 20 --test-data kodak --corruption 14 --severity 3 --psd

To get the Fourier heatmaps, we need to test model on all Fourier-shifted versions datasets. We use Fourier heatmaps of size 64 x 64. Because of their symmetry, we need to test the model on 33 * 64 = 2112 of these shifts. To test a particular shift, use the --fhm-idx flag as shown below:

for i in {0..2111}
do
    python test.py --run-dir runs/sh_fr_dense/lmbda-0.05 --resume runs/sh_fr_dense/lmbda-0.05/checkpoints/epoch_4999.state --fhm-idx $i
done

Visualizing D, G, R and Fourier Heatmaps

Code for visualizing PSDS ($\mathcal{D}, \mathcal{G}, \mathcal{R}$) and Fourier Heatmaps is in visualize/visualize.py. An example of how to use these is in visualize_example.py.

Citation

If you use our work, please cite:

@inproceedings{
lieberman2023neural,
title={Neural Image Compression: Generalization, Robustness, and Spectral Biases},
author={Kelsey Lieberman and James Diffenderfer and Charles Godfrey and Bhavya Kailkhura},
booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
year={2023},
url={https://openreview.net/forum?id=FxRfAIj4s2}
}

License

This repository is distributed under the terms of the MIT license. LLNL-CODE-856963.