Skip to content

PyTorch implementation of Full Resolution Image Compression with Recurrent Neural Networks

Notifications You must be signed in to change notification settings

1zb/pytorch-image-comp-rnn

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Apr 19, 2018
d7150f2 · Apr 19, 2018

History

23 Commits
Mar 2, 2018
Nov 1, 2017
Nov 2, 2017
Sep 6, 2017
Nov 2, 2017
Sep 6, 2017
Sep 6, 2017
Sep 6, 2017
Sep 1, 2017
Nov 1, 2017
Nov 1, 2017
Sep 6, 2017
Apr 18, 2018
Nov 1, 2017
Nov 2, 2017
Sep 5, 2017

Repository files navigation

Full Resolution Image Compression with Recurrent Neural Networks

https://arxiv.org/abs/1608.05148v2

Requirements

  • PyTorch 0.2.0

Train

python train.py -f /path/to/your/images/folder/like/mscoco

Encode and Decode

Encode

python encoder.py --model checkpoint/encoder_epoch_00000005.pth --input /path/to/your/example.png --cuda --output ex --iterations 16

This will output binary codes saved in .npz format.

Decode

python decoder.py --model checkpoint/encoder_epoch_00000005.pth --input /path/to/your/example.npz --cuda --output /path/to/output/folder

This will output images of different quality levels.

Test

Get Kodak dataset

bash test/get_kodak.sh

Encode and decode with RNN model

bash test/enc_dec.sh

Encode and decode with JPEG (use convert from ImageMagick)

bash test/jpeg.sh

Calculate SSIM

bash test/calc_ssim.sh

Draw rate-distortion curve

python test/draw_rd.py

Result

LSTM (Additive Reconstruction), before entropy coding

Rate-distortion

Rate-distortion

kodim10.png

Original Image

Original Image

Below Left: LSTM, SSIM=0.865, bpp=0.125

Below Right: JPEG, SSIM=0.827, bpp=0.133

bpp-0.125-0.133-ssim-0.865-0.827

Below Left: LSTM, SSIM=0.937, bpp=0.250

Below Right: JPEG, SSIM=0.918, bpp=0.249

bpp-0.250-0.249-ssim-0.937-0.918

Below Left: LSTM, SSIM=0.963, bpp=0.375

Below Right: JPEG, SSIM=0.951, bpp=0.381

bpp-0.375-0.381-ssim-0.963-0.951

What's inside

  • train.py: Main program for training.
  • encoder.py and decoder.py: Encoder and decoder.
  • dataset.py: Utils for reading images.
  • metric.py: Functions for Calculatnig MS-SSIM and PSNR.
  • network.py: Modules of encoder and decoder.
  • modules/conv_rnn.py: ConvLSTM module.
  • functions/sign.py: Forward and backward for binary quantization.

Official Repo

https://github.com/tensorflow/models/tree/master/compression