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.
- Scale-hyperprior
- ELIC
- A variable-rate version of scale-hyperprior trained using loss conditional training.
- A sparse version of scale-hyperprior pruned using gradual magnitude pruning.
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.
- Image compressor options include NIC models listed above or classic codecs (i.e., JPEG, JPEG2000, and VTM (equivalently VVC)).
- If testing a NIC model, use
--resume
to specify a checkpoint for testing.- If testing a variable-rate NIC model, use
--variable-rate
and--eval-lambda
to test the model on a particular lambda.
- If testing a variable-rate NIC model, use
- If testing a classic codec, use
--classic-codec
and-q
to specify which codec and quality. - If testing without compression, use
--no-compression
.
- If testing a NIC model, use
- Dataset choices are CLIC, Kodak, or ImageNet (validation split) and their -C versions.
- Use
--test-dataset
to specify which dataset to use. - To test the clean dataset (e.g., CLIC), set
--corruption
and--severity
to 0. - To test a corruption from the -C dataset, set
--corruption
to the number 1-15 corresponding to the desired corruption (the order of corruptions can be found indata/corruptions.py
) and--severity
to the number 1-5 corresponding to the desired severity. - To test a Fourier-shifted dataset, use
--fhm-idx
.
- Use
- By default,
test.py
will save the rate-distortion (bpp/PSNR) metrics in a json file. Additional arguments for more evaluation functionality include:--psd
: save the arrays of the PSDs for the model on that dataset. If the corruption is clean (i.e., corruption=0), then one PSD ($\mathcal{D}$) will be saved. If the corruption is not clean (i.e., using OOD data), then two PSDs ($\mathcal{G}$ and $\mathcal{R}$) will be saved.--save-images
: saves reconstructed images.--classify
: use a pre-trained ResNet model to classify the images after compression. Can only be used with ImageNet dataset.
- See arguments in
args/test_args.py
for more details.
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.