| --- |
| 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:danjacobellis@utexas.edu) |
|
|
| ## 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 |
| ``` |
|
|
|
|
|
|
|
|
| |
|  |
