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.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.