CompressAI

CompressAI（发音为compress-ay）是一个用于端到端压缩研究的PyTorch库和评估平台。

CompressAI目前提供：

用于基于深度学习的数据压缩的自定义操作、层和模型
官方TensorFlow压缩库的部分移植
用于学习图像压缩的预训练端到端压缩模型
评估脚本，用于将学习模型与经典图像/视频压缩编解码器进行比较

Kodak数据集上的PSNR性能图

注意：多GPU支持目前处于实验阶段。

安装

CompressAI支持Python 3.8+和PyTorch 1.7+。

pip安装：

pip install compressai

注意：Linux和MacOS提供预编译的wheel包。

从源代码安装：

还需要C++17编译器、最新版本的pip（19.0+）以及常见的Python包（完整列表请参见setup.py）。

要在本地开始使用并安装CompressAI的开发版本，请在虚拟环境中运行以下命令：

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

对于自定义安装，您还可以运行以下命令之一：

pip install -e '.[dev]'：安装开发所需的包（测试、代码检查、文档）
pip install -e '.[tutorials]'：安装教程所需的包（notebooks）
pip install -e '.[all]'：安装所有可选包

注意：未来将发布Docker镜像。不官方支持Conda环境。

文档

使用方法

示例

脚本和notebook示例可以在examples/目录中找到。

要使用提供的预训练模型对图像进行编码/解码，请运行codec.py示例：

python3 examples/codec.py --help

examples/train.py中提供了一个使用率失真损失的示例训练脚本。您可以用CompressAI中实现的自定义模型替换训练脚本中使用的模型，然后运行脚本以获得简单的训练流程：

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

**注意：**训练示例使用自定义的ImageFolder结构。

examples目录中还提供了一个Jupyter notebook，展示了如何使用预训练模型进行学习图像压缩：

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

评估

要在自己的数据集上评估训练好的模型，CompressAI提供了一个评估脚本：

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

要评估提供的预训练模型：

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

要绘制bench/eval_model模拟的结果（默认需要matplotlib）：

python3 -m compressai.utils.plot --help

要评估传统编解码器：

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

对于视频，可以运行类似的测试，CompressAI目前仅包括ssf2020：

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

测试

使用pytest运行测试：

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

可以使用-m "not slow"选项跳过慢速测试。

许可证

CompressAI采用BSD 3-Clause Clear许可证

贡献

我们欢迎反馈和贡献。如果发现错误、请求增强功能或有任何问题，请开启GitHub issue。

在贡献之前，请阅读CONTRIBUTING.md文件。

作者

Jean Bégaint、Fabien Racapé、Simon Feltman和Hyomin Choi，InterDigital AI实验室。

引用

如果您使用本项目，请引用相关模型和数据集的原始出版物，并引用本项目如下：

@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},
}

对于与可变比特率模型相关的任何工作，请引用

@article{kamisli2024dcc_vbrlic,
	title={Variable-Rate Learned Image Compression with Multi-Objective Optimization and Quantization-Reconstruction Offsets},
	author={Kamisli, Fatih and Racap{\'e}, Fabien and Choi, Hyomin},
	year={2024},
	booktitle={2024 Data Compression Conference (DCC)},
	eprint={2402.18930},
}