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.

baoyu2020/TinyLIC