魔鬼藏在细节中:基于窗口注意力的图像压缩

PyTorch 实现论文"魔鬼藏在细节中:基于窗口注意力的图像压缩"。CVPR2022。本仓库基于 CompressAI。我们保留了训练和评估的脚本，移除了其他组件。主要更改在 compressai/models 中。关于官方代码发布,请参见 CompressAI。

关于

本仓库定义了"魔鬼藏在细节中:基于窗口注意力的图像压缩"中用于学习图像压缩的基于 CNN 的模型和基于 Transformer 的模型。

cnn_架构

基于 CNN 的模型架构。

stf_架构

基于 Transformer 的模型 (STF) 架构。

安装

安装 CompressAI 和开发所需的包。

conda create -n compress python=3.7
conda activate compress
pip install compressai
pip install pybind11
git clone https://github.com/Googolxx/STF stf
cd stf
pip install -e .
pip install -e '.[dev]'

注意: Linux 和 MacOS 提供了 wheel 包。

使用方法

训练

在 train.py 中提供了一个使用率失真损失的示例训练脚本。

训练基于 CNN 的模型:

CUDA_VISIBLE_DEVICES=0,1 python train.py -d /path/to/image/dataset/ -e 1000 --batch-size 16 --save --save_path /path/to/save/ -m cnn --cuda --lambda 0.0035
例如, CUDA_VISIBLE_DEVICES=0,1 python train.py -d openimages -e 1000 --batch-size 16 --save --save_path ckpt/cnn_0035.pth.tar -m cnn --cuda --lambda 0.0035

训练基于 Transformer 的模型(STF):

CUDA_VISIBLE_DEVICES=0,1 python train.py -d /path/to/image/dataset/ -e 1000 --batch-size 16 --save --save_path /path/to/save/ -m stf --cuda --lambda 0.0035

评估

要在您自己的数据集上评估训练好的模型,评估脚本为:

CUDA_VISIBLE_DEVICES=0 python -m compressai.utils.eval_model -d /path/to/image/folder/ -r /path/to/reconstruction/folder/ -a stf -p /path/to/checkpoint/ --cuda

CUDA_VISIBLE_DEVICES=0 python -m compressai.utils.eval_model -d /path/to/image/folder/ -r /path/to/reconstruction/folder/ -a cnn -p /path/to/checkpoint/ --cuda

数据集

下载 OpenImages 的脚本在 downloader_openimages.py 中提供。请先安装 fiftyone。

结果

可视化

可视化01

重建图像 kodim01.png 的可视化。

可视化07

重建图像 kodim07.png 的可视化。

率失真曲线

kodak_rd

Kodak 上的率失真曲线。

clic_rd

CLIC 专业验证数据集上的率失真曲线。

Kodak上的编解码器效率

方法	编码(秒)	解码(秒)	PSNR	bpp
CNN	0.12	0.12	35.91	0.650
STF	0.15	0.15	35.82	0.651

预训练模型

使用从OpenImages数据集中随机选择的30万张图像从头开始训练的预训练模型（针对MSE优化）。

方法	Lambda	链接
CNN	0.0018	cnn_0018
CNN	0.0035	cnn_0035
CNN	0.0067	cnn_0067
CNN	0.025	cnn_025
STF	0.0018	stf_0018
STF	0.0035	stf_0035
STF	0.0067	stf_0067
STF	0.013	stf_013
STF	0.025	stf_025
STF	0.0483	stf_0483

其他预训练模型将陆续发布。

引用

@inproceedings{zou2022the,
  title={The Devil Is in the Details: Window-based Attention for Image Compression},
  author={Zou, Renjie and Song, Chunfeng and Zhang, Zhaoxiang},
  booktitle={CVPR},
  year={2022}
}