Combiner

Offical implementation of Compression with Bayesian Implicit Neural Representations, NeurIPS 2023 spotlight.

And we are excited to have a follow-up powerful coding scheme accepted to ICLR2024, which can outperform autoencoder-based codecs on low-resolution images at low bitrates:

RECOMBINER: Robust and Enhanced Compression with Bayesian Implicit Neural Representations

Installation

We recommend using conda environment for installation.

conda create --name $ENV_NAME
conda activate $ENV_NAME
conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c pytorch
pip3 install -r requirements.txt

Preparing training and test data.

Currently, we only provide the code for Kodak dataset.

We recommend using cropped patches as training set to learn the model prior. The training datasets can be generated from CLIC dataset or DIV2K dataset. The folder structure of training set is like this:

|-- trainset
	|-- img0.png
	|-- img1.png
	|-- img2.png
	...

To align with the resolution of Kodak, the training images should be randomly cropped into 512x768 or 768x512 patches.

Running the code

Run the training script train_model_prior.py to learn the model prior by coordinate descent (parallel training):

CUDA_VISIBLE_DEVICES=0,1,2,3 python -u train_model_prior.py -c=./cfgs/model_prior_kodak.yaml 

Run the encoding scripts for a test image on Kodak dataset. First, learn the model posterior of all the 24 Kodak dataset (parallel optimization).

CUDA_VISIBLE_DEVICES=0,1,2,3 python -u encode_learn_posterior.py -c=./cfgs/model_posterior_kodak.yaml 

Specifying a test image, finetune the model posterior progressively and encode the model parameters into binary file.

CUDA_VISIBLE_DEVICES=0 python -u encode_tune_posterior.py -c=./cfgs/model_posterior_kodak.yaml --image_id=3

You can decode the binary file by running

CUDA_VISIBLE_DEVICES=0 python -u decode_posterior.py -c=./cfgs/model_posterior_kodak.yaml --image_id=3

Configuration

You can adjust the hyper parameters of training and test settings by modifying the files in ./cfgs

Citation

If you use this library for research purposes, please cite:

@inproceedings{combiner,
  author       = {Zongyu Guo and
                  Gergely Flamich and
                  Jiajun He and
                  Zhibo Chen and
                  Jos{\'{e}} Miguel Hern{\'{a}}ndez{-}Lobato},
  title        = {Compression with Bayesian Implicit Neural Representations},
  booktitle    = {Advances in Neural Information Processing Systems 36: Annual Conference
                  on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans,
                  LA, USA, December 10 - 16, 2023},
  year         = {2023}
}