README.md

metadata

datasets:
  - danjacobellis/imagenet_hq

Lightweight Learned Image Compression (LLIC)

Installation

1. Follow the installation instructions for torch and compressai
2. Install LLIC via pip: pip install LLIC

Pre-trained checkpoints

An imagenet-trained checkpoint for RGB images is available on huggingface: LLIC_rgb_v0.0.1.pth

Request access to other checkpoints (grayscale, hyperspectral, microscopy, etc)

Usage example

import torch
import zlib
import numpy as np
import compressai
from io import BytesIO
from IPython.display import display
from PIL import Image
from LLIC import LLIC
from torchvision.transforms import ToPILImage, PILToTensor

Load the model

checkpoint = torch.load("LLIC_rgb_v0.0.1.pth",map_location="cpu")
codec = LLIC.RateDistortionAutoEncoder()
codec.load_state_dict(checkpoint['model_state_dict'])

<All keys matched successfully>

Download example image

!wget https://r0k.us/graphics/kodak/kodak/kodim05.png

original_image = Image.open("kodim05.png")
original_image

The analysis and synthesis transforms expect dimensions to be multiples of of 16. Zero padding can be applied otherwise.

def pad(x, p=2**5):
    h, w = x.size(2), x.size(3)
    pad, _ = compressai.ops.compute_padding(h, w, min_div=p)
    return torch.nn.functional.pad(x, pad, mode="constant", value=0)

def preprocess(pil_image):
    tensor = PILToTensor()(pil_image)
    tensor = tensor.unsqueeze(0)
    tensor = tensor.to(torch.float)
    tensor = tensor/255
    tensor = tensor - 0.5
    return pad(tensor)

Compress the image and save file

padded_image = preprocess(original_image)
original_size = padded_image.shape
compressed_image, compressed_shape = LLIC.compress(padded_image, codec)
with open("kodim05.llic", 'wb') as f:
    f.write(compressed_image)

Decompress and view the image

def crop(x, size):
    H, W = x.size(2), x.size(3)
    h, w = size
    _, unpad = compressai.ops.compute_padding(h, w, out_h=H, out_w=W)
    return torch.nn.functional.pad(x, unpad, mode="constant", value=0)

def postprocess(tensor):
    tensor = tensor[0] + 0.5
    tensor = 255*tensor
    tensor = tensor.clamp(0,255)
    tensor = tensor.to(torch.uint8)
    pil_image = ToPILImage()(tensor)
    return pil_image

with open("kodim05.llic", 'rb') as f:
    compressed_image = f.read()
    tensor = LLIC.decompress(compressed_image, compressed_shape, codec)
    recovered_image = postprocess(crop(tensor, (512,768)))

recovered_image