High-Efficiency Lossy Image Coding Through Adaptive Neighborhood Information Aggregation
Pytorch Implementation of our paper "High-Efficiency Lossy Image Coding Through Adaptive Neighborhood Information Aggregation"[arXiv].
More details can be found at the homepage.
Installation
To get started locally and install the development version of our work, run the following commands (The docker environment is recommended):
git clone https://github.com/lumingzzz/TinyLIC.git
cd TinyLIC
pip install -U pip && pip install -e .
pip install timm
Usage
Train
We use the Flicker2W dataset for training, and the script for preprocessing.
Run the script for a simple training pipeline:
python examples/train.py -m tinylic -d /path/to/my/image/dataset/ --epochs 400 -lr 1e-4 --batch-size 8 --cuda --save
The training checkpoints will be generated in the "pretrained" folder at current directory. You can change the default folder by modifying the function "init" in "./expample/train.py".
Evaluation
Pre-trained models for Quality 3 (lambda = 0.0067) can be downloaded in NJU Box. Since the codes are reorgnized and retrained using the newest version (1.2.0) of CompressAI, the performance may be slightly different from that in the paper.
The mse optimized results can be found in /results for reference.
An example to evaluate model:
python -m compressai.utils.eval_model checkpoint path/to/eval/data/ -a tinylic -p path/to/pretrained/model --cuda
The simulation results in our paper are all with "--entropy-estimation" enabled for quick evaluation. There may be a slight deviation against the actual ''compress-decompress'' while the contrast relationships are consistent.
Citation
If you find this work useful for your research, please cite:
@article{lu2022high,
title={High-Efficiency Lossy Image Coding Through Adaptive Neighborhood Information Aggregation},
author={Lu, Ming and Ma, Zhan},
journal={arXiv preprint arXiv:2204.11448},
year={2022}
}
Acknowledgement
The framework is based on CompressAI, we add our modifications in compressai.models.tinylic and compressai.layers for usage.
The TinyLIC model is partially built upon the Swin Transformer and the open sourced unofficial implementation of checkerboard shaped context model. We thank the authors for sharing their code.
Contact
If you have any question, please contact me via luming@smail.nju.edu.cn.