danjacobellis/imagenet_RDAE

imagenet-RDAE数据集是基于ImageNet-1k数据集，通过Rate-Distortion Autoencoder（RDAE）压缩至7GB。RDAE是在Vimeo90k数据集上训练的，并提供了训练脚本的链接。该数据集用于图像压缩和解压缩，文件中包含了如何使用预训练模型进行编码和解码的示例代码，以及如何比较压缩图像与JPEG压缩图像的质量。

原始信息汇总

imagenet-RDAE

该数据集包含使用Rate-Distortion Autoencoder压缩的ImageNet-1k Dataset，原始大小为216 GB，压缩后为7GB。

使用RDAE

RDAE在Vimeo90k上进行训练，训练脚本可在此处获取：训练脚本。

示例代码

以下是使用预训练模型进行编码和解码的示例：

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

获取示例图像

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

图像预处理

自动编码器期望图像范围为$[-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

网络结构

自动编码器使用三层卷积和generalized divisive normalization：

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

压缩过程

压缩分为两个阶段：有损分析变换和无损熵编码。使用python zlib模块进行熵编码：

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) ẑ = np.frombuffer(decompressed, dtype=np.int8) ẑ = ẑ.reshape(original_shape) return ẑ

def synthesis_transform(ẑ): ẑ = torch.tensor(ẑ).to("cuda").to(torch.float) x̂ = net.decode(ẑ).detach().to("cpu") return x̂

def decompress(compressed_img, original_shape): ẑ = entropy_decoder(compressed_img,original_shape) x̂ = synthesis_transform(ẑ) return x̂

加载预训练网络

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

压缩和解压图像

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

比较压缩质量

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

查看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);