arxiv:1511.06281 region:us
# imagenet-RDAE

This dataset consists of the [ImageNet-1k Dataset (original size 216 GB)](https://huggingface.co/datasets/imagenet-1k) compressed to 7GB using a [Rate-Distortion Autoencoder](https://www.cns.nyu.edu/~lcv/iclr2017/).

## Using the RDAE

The rate-distortion autoencoder was trained on [Vimeo90k](http://toflow.csail.mit.edu/). The [training script is available here](https://github.com/danjacobellis/moco/blob/main/train_RDAE_on_vimeo.ipynb).

An example showing how to use the pretrained model for both encoding and decoding is provided below.


```python
import os
import zlib
import requests
import numpy as np
import PIL.Image as Image
from IPython.display import display

import torch
import torch.nn as nn
from torchvision import transforms

import compressai
from compressai.entropy_models import EntropyBottleneck
from compressai.layers import GDN
from compressai.models import CompressionModel
from compressai.models.utils import conv, deconv
```

### Get an example image


```python
r = requests.get("https://r0k.us/graphics/kodak/kodak/kodim05.png", stream=True)
img = Image.open(r.raw); img
```




    
![png](README_files/README_4_0.png)
    



The autoencoder expects images to be scaled to the range $[-1.0,1.0]$.


```python
def pil_to_pt(img):
    t = transforms.functional.pil_to_tensor(img)
    t = t.to(torch.float)
    t = t/255
    t = t-0.5
    t = t.unsqueeze(0)
    return t
def pt_to_pil(t):
    t = t+0.5
    t = t*255
    t = torch.clamp(t, min=-0.49, max=255.49)
    t = t.round()
    t = t.to(torch.uint8)
    return t
```

The autoencoder uses three convolutional layers and [generalized divisive normalization](https://arxiv.org/abs/1511.06281).


```python
class Network(CompressionModel):
    def __init__(self, N=128):
        super().__init__()
        self.entropy_bottleneck = EntropyBottleneck(N)
        self.encode = nn.Sequential(
            conv(3, N),
            GDN(N),
            conv(N, N),
            GDN(N),
            conv(N, N),
        )

        self.decode = nn.Sequential(
            deconv(N, N),
            GDN(N, inverse=True),
            deconv(N, N),
            GDN(N, inverse=True),
            deconv(N, 3),
        )

    def forward(self, x):
        y = self.encode(x)
        y_hat, y_likelihoods = self.entropy_bottleneck(y)
        x_hat = self.decode(y_hat)
        return x_hat, y_likelihoods
```

Compression occurs in two stages: A lossy analysis transform and lossless entropy coding. The python zlib module is used for entropy coding.


```python
def lossy_analysis_transform(img):
    x = pil_to_pt(img).to("cuda")
    z = net.encode(x).round().to(torch.int8).detach().to("cpu").numpy()
    return z
    
def lossless_entropy_encode(z):
    original_shape = z.shape
    compressed_img = zlib.compress(z.tobytes(), level=9)
    return compressed_img, original_shape

def compress(img):
    z = lossy_analysis_transform(img)
    compressed_img, original_shape = lossless_entropy_encode(z)
    return compressed_img, original_shape
```


```python
def entropy_decoder(compressed_img,original_shape):
    decompressed = zlib.decompress(compressed_img)
    ẑ = np.frombuffer(decompressed, dtype=np.int8)
    ẑ = ẑ.reshape(original_shape)
    return ẑ
    
def synthesis_transform(ẑ):
    ẑ = torch.tensor(ẑ).to("cuda").to(torch.float)
    x̂ = net.decode(ẑ).detach().to("cpu")
    return x̂
    
def decompress(compressed_img, original_shape):
    ẑ = entropy_decoder(compressed_img,original_shape)
    x̂ = synthesis_transform(ẑ)
    return x̂
```

Load the pretrained network.


```python
net = Network()
net = net.to("cuda")
checkpoint = torch.load("checkpoint.pth")
net.load_state_dict(checkpoint['model_state_dict'])
```




    <All keys matched successfully>



Save a copy using JPEG for comparison.


```python
img.save("kodim05.jpg", "JPEG", quality=5)
jpeg = Image.open("kodim05.jpg");
```

Compress the and decompress the image.


```python
compressed_img, original_shape = compress(img);
print("Bytes in compressed image:", len(compressed_img))
x̂ = decompress(compressed_img, original_shape)
```

    Bytes in compressed image: 14391


Compare the quality of the compressed image to JPEG and the original


```python
print('original image');
display(img)
print("JPEG compression ratio", (3*512*768)/os.path.getsize('kodim05.jpg'))
display(jpeg);
print("RDAE compression ratio", (3*512*768)/len(compressed_img))
display(transforms.ToPILImage()(pt_to_pil(x̂)[0]))
```

    original image



    
![png](README_files/README_19_1.png)
    


    JPEG compression ratio 77.25771170345143



    
![png](README_files/README_19_3.png)
    


    RDAE compression ratio 81.97123202001251



    
![png](README_files/README_19_5.png)
    


View the distribution of codes produced by the RDAE.


```python
import matplotlib.pyplot as plt
z = lossy_analysis_transform(img)
histogram = torch.histogram(torch.tensor(z).to(torch.float),bins=33, range=(-16.5,16.5))
x = np.convolve(histogram.bin_edges,[0.5,0.5])[1:-1]
y = np.log(1+histogram.hist)
plt.bar(x,y);
```


    
![png](README_files/README_21_0.png)
    

