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
Note: Multi-GPU support is now experimental.
CompressAI supports python 3.6+ and PyTorch 1.4+.
pip install compressai
Note: wheels are available for Linux and MacOS.
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
Script and notebook examples can be found in the
To encode/decode images with the provided pre-trained models, run the
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
pip install -U ipython jupyter ipywidgets matplotlib jupyter notebook examples/
To evaluate a trained model on your own dataset, CompressAI provides an
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
Run tests with
pytest -sx --cov=compressai --cov-append --cov-report term-missing tests
Slow tests can be skipped with the
-m "not slow" option.
- Jean Bégaint, Fabien Racapé, Simon Feltman and Akshay Pushparaja, InterDigital AI Lab.