[ACMMM 2021 Oral] Enhanced Invertible Encoding for Learned Image Compression

Overview

InvCompress

Official Pytorch Implementation for "Enhanced Invertible Encoding for Learned Image Compression", ACMMM 2021 (Oral)

Figure: Our framework

Acknowledgement

The framework is based on CompressAI, we add our model in compressai.models.ours, compressai.models.our_utils. We modify compressai.utils, compressai.zoo, compressai.layers and examples/train.py for usage. Part of the codes benefit from Invertible-Image-Rescaling.

Introduction

In this paper, we target at structuring a better transformation between the image space and the latent feature space. Instead of employing previous autoencoder style networks to build this transformation, we propose an enhanced Invertible Encoding Network with invertible neural networks (INNs) to largely mitigate the information loss problem for better compression. To solve the challenge of unstable training with INN, we propose an attentive channel squeeze layer to flexibly adjust the feature dimension for a lower bit rate. We also present a feature enhancement module with same-resolution transforms and residual connections to improve the network nonlinear representation capacity.

[Paper]

Figure: Our results

Installation

As mentioned in CompressAI, "A C++17 compiler, a recent version of pip (19.0+), and common python packages are also required (see setup.py for the full list)."

git clone https://github.com/xyq7/InvCompress.git
cd InvCompress/codes/
conda create -n invcomp python=3.7 
conda activate invcomp
pip install -U pip && pip install -e .
conda install -c conda-forge tensorboard

Usage

Evaluation

If you want evaluate with pretrained model, please download from Google drive or Baidu cloud (code: a7jd) and put in ./experiments/

Some evaluation dataset can be downloaded from kodak dataset, CLIC

Note that as mentioned in original CompressAI, "Inference on GPU is not recommended for the autoregressive models (the entropy coder is run sequentially on CPU)." So for inference of our model, please run on CPU.

python -m compressai.utils.eval_model checkpoint $eval_data_dir -a invcompress -exp $exp_name -s $save_dir

An example: to evaluate model of quality 1 optimized with mse on kodak dataset.

python -m compressai.utils.eval_model checkpoint ../data/kodak -a invcompress -exp exp_01_mse_q1 -s ../results/exp_01

If you want to evaluate your trained model on own data, please run update before evaluation. An example:

python -m compressai.utils.update_model -exp $exp_name -a invcompress
python -m compressai.utils.eval_model checkpoint $eval_data_dir -a invcompress -exp $exp_name -s $save_dir

Train

We use the training dataset processed in the repo. We further preprocess with /codes/scripts/flicker_process.py Training setting is detailed in the paper. You can also use your own data for training.

python examples/train.py -exp $exp_name -m invcompress -d $train_data_dir --epochs $epoch_num -lr $lr --batch-size $batch_size --cuda --gpu_id $gpu_id --lambda $lamvda --metrics $metric --save 

An example: to train model of quality 1 optimized with mse metric.

python examples/train.py -exp exp_01_mse_q1 -m invcompress -d ../data/flicker --epochs 600 -lr 1e-4 --batch-size 8 --cuda --gpu_id 0 --lambda 0.0016 --metrics mse --save 

Other usage please refer to the original library CompressAI

Citation

If you find this work useful for your research, please cite:

@inproceedings{xie2021enhanced,
    title = {Enhanced Invertible Encoding for Learned Image Compression}, 
    author = {Yueqi Xie and Ka Leong Cheng and Qifeng Chen},
    booktitle = {Proceedings of the ACM International Conference on Multimedia},
    year = {2021}
}

Contact

Feel free to contact us if there is any question. (YueqiXIE, [email protected]; Ka Leong Cheng, [email protected])

Comments
  • PSNR differences between two different settings.

    PSNR differences between two different settings.

    Hi! Thanks for your great work!

    During measuring PSNR values of test images, is there any difference between them? Since the quantization method of 'forward()' is slightly different from that of 'compress()', I'm wondering how much the difference occurs.

    opened by junhyukk 9
  • Is it convenient to give a pre-trained model under higher compression

    Is it convenient to give a pre-trained model under higher compression

    Thanks for your last reply! When I try to train the model under higher compression (nearing 0.06), I get a poor performance. So do you have the pretrained model under higher compression ?

    opened by JXH-SHU 2
  • Total patches

    Total patches

    Hello,

    Great work and awesome code. Thank you very much for sharing. I see that total 20,745 images are present in Flickr dataset out of which 200 random images are taken away for validation, and then generated patches of size 256*256. If I use the generate patch script, it results huge number of patches. How many patches are present in your training dataset?

    Regards Priyanka

    opened by prmudgal 1
  • Need for speed !

    Need for speed !

    Hi,

    Great work ! I'm currently using it in my phd research project.

    If I understand the code correctly, the context model (ctx_p) is significantly slowing down the entropy coding/decoding computation...

    Is there any way to make this much faster ?

    Thanks.

    opened by jjeremy40 1
  • Fail to load pretrained model

    Fail to load pretrained model

    Hi, I try to load the pre-trained model and fine-tune it in my own dataset. However, there exists a bug in loading the model:

    https://github.com/xyq7/InvCompress/blob/17deb0f0e2f28432e6d5b38d0898ead58ccbcde4/codes/examples/train.py#L444

    In this line, there has an KeyError: 'state_dict'.

    It seems that the model is not saved as train.py in line 501:

    save_checkpoint( { "epoch": epoch + 1, "state_dict": net.state_dict(), "loss": loss, "optimizer": optimizer.state_dict(), "aux_optimizer": aux_optimizer.state_dict(), "lr_scheduler": lr_scheduler.state_dict(), }, is_best, os.path.join('../experiments', args.experiment, 'checkpoints', "checkpoint_%03d.pth.tar" % (epoch + 1)) )

    opened by lep990816 1
  • some problems about the pre-trained model of q2

    some problems about the pre-trained model of q2

    Thank you for your codes. However when I try to use the pre-trained model (exp_02_mse_q2), some errors occurred. Other pre-trained model including (q1, q3,q4,q5) can work normally.

    opened by JXH-SHU 1
  • Small bug of the network architecture

    Small bug of the network architecture

    Hi Yueqi, thanks for your great work!

    I am wondering about a small bug of the network architecture. According to the paper, for the SSIM models, the para N for the first two lambda values is 128, while that for the last three lambda values is 192. However, according to the code, the para N for the first four lambda values is 128, while that for the last four lambda values is 192.

    I want to make sure that if I want to train the all five SSIM models, my training code should be:

    python examples/train.py ../ -exp exp_ssim_q1_01 -m invcompress -d ../data/flickr2w --epochs 600 -lr 1e-4 --batch-size 8 -n 8 --cuda --gpu_id 7 -q 1 --lambda 6 --metrics ms-ssim --save --seed 7
    
    python examples/train.py ../ -exp exp_ssim_q2_01 -m invcompress -d ../data/flickr2w --epochs 600 -lr 1e-4 --batch-size 8 -n 8 --cuda --gpu_id 6 -q 1 --lambda 12 --metrics ms-ssim --save --seed 7
    
    python examples/train.py ../ -exp exp_ssim_q3_01 -m invcompress -d ../data/flickr2w --epochs 600 -lr 1e-4 --batch-size 8 -n 8 --cuda --gpu_id 5 -q 5 --lambda 40 --metrics ms-ssim --save --seed 7
    
    python examples/train.py ../ -exp exp_ssim_q4_01 -m invcompress -d ../data/flickr2w --epochs 600 -lr 1e-4 --batch-size 8 -n 8 --cuda --gpu_id 4 -q 5 --lambda 120 --metrics ms-ssim --save --seed 7
    
    python examples/train.py ../ -exp exp_ssim_q5_01 -m invcompress -d ../data/flickr2w --epochs 600 -lr 1e-4 --batch-size 8 -n 8 --cuda --gpu_id 3 -q 5 --lambda 220 --metrics ms-ssim --save --seed 7
    

    Is it correct? Thank you.

    opened by ryanxingql 0
Owner
null
CoReD: Generalizing Fake Media Detection with Continual Representation using Distillation (ACMMM'21 Oral Paper)

CoReD: Generalizing Fake Media Detection with Continual Representation using Distillation (ACMMM'21 Oral Paper) (Accepted for oral presentation at ACM

Minha Kim 1 Nov 12, 2021
Incremental Transformer Structure Enhanced Image Inpainting with Masking Positional Encoding (CVPR2022)

Incremental Transformer Structure Enhanced Image Inpainting with Masking Positional Encoding by Qiaole Dong*, Chenjie Cao*, Yanwei Fu Paper and Supple

Qiaole Dong 190 Dec 27, 2022
Learned image compression

Overview Pytorch code of our recent work A Unified End-to-End Framework for Efficient Deep Image Compression. We first release the code for Variationa

Jiaheng Liu 163 Dec 4, 2022
A Pytorch Implementation of a continuously rate adjustable learned image compression framework.

GainedVAE A Pytorch Implementation of a continuously rate adjustable learned image compression framework, Gained Variational Autoencoder(GainedVAE). N

null 39 Dec 24, 2022
An Image compression simulator that uses Source Extractor and Monte Carlo methods to examine the post compressive effects different compression algorithms have.

ImageCompressionSimulation An Image compression simulator that uses Source Extractor and Monte Carlo methods to examine the post compressive effects o

James Park 1 Dec 11, 2021
【ACMMM 2021】DSANet: Dynamic Segment Aggregation Network for Video-Level Representation Learning

DSANet: Dynamic Segment Aggregation Network for Video-Level Representation Learning (ACMMM 2021) Overview We release the code of the DSANet (Dynamic S

Wenhao Wu 46 Dec 27, 2022
Revisiting Discriminator in GAN Compression: A Generator-discriminator Cooperative Compression Scheme (NeurIPS2021)

Revisiting Discriminator in GAN Compression: A Generator-discriminator Cooperative Compression Scheme (NeurIPS2021) Overview Prerequisites Linux Pytho

Shaojie Li 34 Mar 31, 2022
[CVPR2021] Invertible Image Signal Processing

Invertible Image Signal Processing This repository includes official codes for "Invertible Image Signal Processing (CVPR2021)". Figure: Our framework

Yazhou XING 281 Dec 31, 2022
Code for "Neural Parts: Learning Expressive 3D Shape Abstractions with Invertible Neural Networks", CVPR 2021

Neural Parts: Learning Expressive 3D Shape Abstractions with Invertible Neural Networks This repository contains the code that accompanies our CVPR 20

Despoina Paschalidou 161 Dec 20, 2022
This is the repository for paper NEEDLE: Towards Non-invertible Backdoor Attack to Deep Learning Models.

This is the repository for paper NEEDLE: Towards Non-invertible Backdoor Attack to Deep Learning Models.

null 1 Oct 25, 2021
InvTorch: memory-efficient models with invertible functions

InvTorch: Memory-Efficient Invertible Functions This module extends the functionality of torch.utils.checkpoint.checkpoint to work with invertible fun

Modar M. Alfadly 12 May 12, 2022
Multi-Scale Vision Longformer: A New Vision Transformer for High-Resolution Image Encoding

Vision Longformer This project provides the source code for the vision longformer paper. Multi-Scale Vision Longformer: A New Vision Transformer for H

Microsoft 209 Dec 30, 2022
The undersampled DWI image using Slice-Interleaved Diffusion Encoding (SIDE) method can be reconstructed by the UNet network.

UNet-SIDE The undersampled DWI image using Slice-Interleaved Diffusion Encoding (SIDE) method can be reconstructed by the UNet network. For Super Reso

TIANTIAN XU 1 Jan 13, 2022
git《FSCE: Few-Shot Object Detection via Contrastive Proposal Encoding》(CVPR 2021) GitHub: [fig8]

FSCE: Few-Shot Object Detection via Contrastive Proposal Encoding (CVPR 2021) This repo contains the implementation of our state-of-the-art fewshot ob

null 233 Dec 29, 2022
Official code release for "Learned Spatial Representations for Few-shot Talking-Head Synthesis" ICCV 2021

Official code release for "Learned Spatial Representations for Few-shot Talking-Head Synthesis" ICCV 2021

Moustafa Meshry 16 Oct 5, 2022
Implementation of Convolutional enhanced image Transformer

CeiT : Convolutional enhanced image Transformer This is an unofficial PyTorch implementation of Incorporating Convolution Designs into Visual Transfor

Rishikesh (ऋषिकेश) 82 Dec 13, 2022
EDPN: Enhanced Deep Pyramid Network for Blurry Image Restoration

EDPN: Enhanced Deep Pyramid Network for Blurry Image Restoration Ruikang Xu, Zeyu Xiao, Jie Huang, Yueyi Zhang, Zhiwei Xiong. EDPN: Enhanced Deep Pyra

null 69 Dec 15, 2022
PyTorch version of the paper 'Enhanced Deep Residual Networks for Single Image Super-Resolution' (CVPRW 2017)

About PyTorch 1.2.0 Now the master branch supports PyTorch 1.2.0 by default. Due to the serious version problem (especially torch.utils.data.dataloade

Sanghyun Son 2.1k Jan 1, 2023
Torch implementation of "Enhanced Deep Residual Networks for Single Image Super-Resolution"

NTIRE2017 Super-resolution Challenge: SNU_CVLab Introduction This is our project repository for CVPR 2017 Workshop (2nd NTIRE). We, Team SNU_CVLab, (B

Bee Lim 625 Dec 30, 2022