CompressAI: a PyTorch library and evaluation platform for end-to-end compression research

TL;DR

CompressAI is a comprehensive PyTorch-based platform that facilitates research, development, and evaluation of end-to-end image and video compression models, including pre-trained models and comparison tools.

Contribution

It introduces a versatile platform with reimplemented state-of-the-art models, evaluation tools, and a focus on end-to-end compression research in PyTorch.

Findings

Reimplemented multiple state-of-the-art models in PyTorch

Provided objective comparison metrics on Kodak dataset

Established a foundation for future video compression research

Abstract

This paper presents CompressAI, a platform that provides custom operations, layers, models and tools to research, develop and evaluate end-to-end image and video compression codecs. In particular, CompressAI includes pre-trained models and evaluation tools to compare learned methods with traditional codecs. Multiple models from the state-of-the-art on learned end-to-end compression have thus been reimplemented in PyTorch and trained from scratch. We also report objective comparison results using PSNR and MS-SSIM metrics vs. bit-rate, using the Kodak image dataset as test set. Although this framework currently implements models for still-picture compression, it is intended to be soon extended to the video compression domain.