File size: 2,911 Bytes
a56ec4d 3fb5bd1 a56ec4d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
---
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
```
 |