|
--- |
|
datasets: |
|
- danjacobellis/imagenet_hq |
|
--- |
|
# Lightweight Learned Image Compression (LLIC) |
|
|
|
## Installation |
|
|
|
1. Follow the installation instructions for [torch](https://pytorch.org/get-started/locally/) and [compressai](https://interdigitalinc.github.io/CompressAI/installation.html) |
|
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](https://huggingface.co/danjacobellis/LLIC/resolve/main/LLIC_rgb_v0.0.1.pth) |
|
|
|
[Request access to other checkpoints (grayscale, hyperspectral, microscopy, etc)](mailto:[email protected]) |
|
|
|
## Usage example |
|
|
|
|
|
```python |
|
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 |
|
|
|
|
|
```python |
|
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 |
|
|
|
|
|
```python |
|
!wget https://r0k.us/graphics/kodak/kodak/kodim05.png |
|
``` |
|
|
|
|
|
```python |
|
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. |
|
|
|
|
|
```python |
|
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 |
|
|
|
|
|
```python |
|
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 |
|
|
|
|
|
```python |
|
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 |
|
``` |
|
|
|
|
|
```python |
|
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))) |
|
``` |
|
|
|
|
|
```python |
|
recovered_image |
|
``` |
|
|
|
|
|
|
|
|
|
|
|
 |
