CompressAI

CompressAI (compress-ay) is a PyTorch library and evaluation platform for end-to-end compression research.

CompressAI currently provides:

custom operations, layers and models for deep learning based data compression
a partial port of the official TensorFlow compression library
pre-trained end-to-end compression models for learned image compression
evaluation scripts to compare learned models against classical image/video compression codecs

Note: Multi-GPU support is now experimental.

Installation

CompressAI supports python 3.6+ and PyTorch 1.7+.

pip:

pip install compressai

Note: wheels are available for Linux and MacOS.

From source:

A C++17 compiler, a recent version of pip (19.0+), and common python packages are also required (see setup.py for the full list).

To get started locally and install the development version of CompressAI, run the following commands in a virtual environment:

git clone https://github.com/InterDigitalInc/CompressAI compressai
cd compressai
pip install -U pip && pip install -e .

For a custom installation, you can also run one of the following commands:

pip install -e '.[dev]': install the packages required for development (testing, linting, docs)
pip install -e '.[tutorials]': install the packages required for the tutorials (notebooks)
pip install -e '.[all]': install all the optional packages

Note: Docker images will be released in the future. Conda environments are not officially supported.

Documentation

Usage

Examples

Script and notebook examples can be found in the examples/ directory.

To encode/decode images with the provided pre-trained models, run the codec.py example:

python3 examples/codec.py --help

An examplary training script with a rate-distortion loss is provided in examples/train.py. You can replace the model used in the training script with your own model implemented within CompressAI, and then run the script for a simple training pipeline:

python3 examples/train.py -d /path/to/my/image/dataset/ --epochs 300 -lr 1e-4 --batch-size 16 --cuda --save

Note: the training example uses a custom ImageFolder structure.

A jupyter notebook illustrating the usage of a pre-trained model for learned image compression is also provided in the examples directory:

pip install -U ipython jupyter ipywidgets matplotlib
jupyter notebook examples/

Evaluation

To evaluate a trained model on your own dataset, CompressAI provides an evaluation script:

python3 -m compressai.utils.eval_model checkpoint /path/to/images/folder/ -a $ARCH -p $MODEL_CHECKPOINT...

To evaluate provided pre-trained models:

python3 -m compressai.utils.eval_model pretrained /path/to/images/folder/ -a $ARCH -q $QUALITY_LEVELS...

To plot results from bench/eval_model simulations (requires matplotlib by default):

python3 -m compressai.utils.plot --help

To evaluate traditional codecs:

python3 -m compressai.utils.bench --help
python3 -m compressai.utils.bench bpg --help
python3 -m compressai.utils.bench vtm --help

For video, similar tests can be run, CompressAI only includes ssf2020 for now:

python3 -m compressai.utils.video.eval_model checkpoint /path/to/video/folder/ -a ssf2020 -p $MODEL_CHECKPOINT...
python3 -m compressai.utils.video.eval_model pretrained /path/to/video/folder/ -a ssf2020 -q $QUALITY_LEVELS...
python3 -m compressai.utils.video.bench x265 --help
python3 -m compressai.utils.video.bench VTM --help
python3 -m compressai.utils.video.plot --help

Tests

Run tests with pytest:

pytest -sx --cov=compressai --cov-append --cov-report term-missing tests

Slow tests can be skipped with the -m "not slow" option.

License

CompressAI is licensed under the BSD 3-Clause Clear License

Contributing

We welcome feedback and contributions. Please open a GitHub issue to report bugs, request enhancements or if you have any questions.

Before contributing, please read the CONTRIBUTING.md file.

Authors

Jean Bégaint, Fabien Racapé, Simon Feltman and Hyomin Choi, InterDigital AI Lab.

Citation

If you use this project, please cite the relevant original publications for the models and datasets, and cite this project as:

@article{begaint2020compressai,
	title={CompressAI: a PyTorch library and evaluation platform for end-to-end compression research},
	author={B{\'e}gaint, Jean and Racap{\'e}, Fabien and Feltman, Simon and Pushparaja, Akshay},
	year={2020},
	journal={arXiv preprint arXiv:2011.03029},
}

Maryamsana-1998/CompressAI