魔鬼藏在细节中:基于窗口注意力的图像压缩
PyTorch 实现论文"魔鬼藏在细节中:基于窗口注意力的图像压缩"。CVPR2022。
本仓库基于 CompressAI。我们保留了训练和评估的脚本,移除了其他组件。主要更改在 compressai/models
中。关于官方代码发布,请参见 CompressAI。
关于
本仓库定义了"魔鬼藏在细节中:基于窗口注意力的图像压缩"中用于学习图像压缩的基于 CNN 的模型和基于 Transformer 的模型。
基于 CNN 的模型架构。
基于 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。
结果
可视化
重建图像 kodim01.png 的可视化。
重建图像 kodim07.png 的可视化。
率失真曲线
Kodak 上的率失真曲线。
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}
}
相关链接
- CompressAI: https://github.com/InterDigitalInc/CompressAI
- Swin-Transformer: https://github.com/microsoft/Swin-Transformer
- Ballé等人的Tensorflow压缩库: https://github.com/tensorflow/compression
- Fabian 'ryg' Giesen的范围不对称数字系统代码: https://github.com/rygorous/ryg_rans
- Kodak图像数据集: http://r0k.us/graphics/kodak/
- Open Images数据集: https://github.com/openimages
- fiftyone: https://github.com/voxel51/fiftyone
- CLIC: https://www.compression.cc/