This repository contains the authors' implementation of several deep learning-based methods related to lossy image compression.
Progressive coding: our models learn a deep hierarchy of latent variables and compress/decompress images in a coarse-to-fine fashion. This feature comes from the hierarchical nature of ResNet VAEs.
Compression performance: our models are powerful in terms of both rate-distortion and decoding speed. Please see the results section below.
Notes on metric computation:
Time is the latency to encode/decode a 512x768 image, averaged over 24 Kodak images. Tested in plain PyTorch (v1.13 + CUDA 11.7) code, ie, no mixed-precision, torchscript, ONNX/TensorRT, etc. \ CPU is Intel 10700k. \ *BD-rate is w.r.t. VTM 18.0, averaged on three common test sets (Kodak, Tecnick TESTIMAGES, and CLIC 2022 test set).
Requirements:
conda install tqdm
Download and Install:
lossy-vae/lvae/paths.py
.cd /pasth/to/lossy-vae
python -m pip install -e .
from lvae import get_model
model = get_model('qarv_base', pretrained=True) # weights are downloaded automatically
model.eval()
model.compress_mode(True) # initialize entropy coding
Encode an image:
model.compress_file('/path/to/image.png', '/path/to/compressed.bits')
Decode an image:
im = model.decompress_file('/path/to/compressed.bits')
# im is a torch.Tensor of shape (1, 3, H, W). RGB. pixel values in [0, 1].
COCO
/path/to/datasets/coco/train2017
lossy-vae/lvae/paths.py
such that
known_datasets['coco-train2017'] = '/path/to/datasets/coco/train2017'
Kodak (link), Tecnick TESTIMAGES (link), and CLIC (link)
python scripts/download-dataset.py --name kodak --datasets_root /path/to/datasets
clic2022-test
tecnick
Then, edit lossy-vae/lvae/paths.py
such that known_datasets['kodak'] = '/path/to/datasets/kodak'
, and similarly for other datasets.
Custom Dataset
lossy-vae/lvae/paths.py
such that known_datasets['custom-name'] = '/path/to/my-dataset'
, where custom-name
is the name of your dataset, and /path/to/my-dataset
is the path to the folder containing images.custom-name
as the dataset name in the training/evaluation scripts.Training and evaluation scripts vary from model to model. For example, qres34m
uses fixed-rate train/eval scheme, while qarv_base
uses variable-rate train/eval scheme. \
Detailed training/evaluation instructions are provided in each model's subfolder (see the section Models).
Code in this repository is freely available for non-commercial use.