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

kodak-psnr

Note: Multi-GPU support is now experimental.

Installation

CompressAI supports python 3.6+ and PyTorch 1.4+.

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 traditional image/video codecs:

python3 -m compressai.utils.bench --help
python3 -m compressai.utils.bench bpg --help
python3 -m compressai.utils.bench vtm --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.

Authors

  • Jean BĂ©gaint, Fabien RacapĂ©, Simon Feltman and Akshay Pushparaja, InterDigital AI Lab.

GitHub

https://github.com/InterDigitalInc/CompressAI