Jacobellis Dan (dgj335)
commited on
Commit
•
3fb5bd1
1
Parent(s):
bd387fd
usage example
Browse files- README.ipynb +0 -0
- README.md +135 -0
- README_files/README_14_0.jpg +3 -0
- README_files/README_14_0.png +3 -0
- README_files/README_6_0.jpg +3 -0
- README_files/README_6_0.png +3 -0
README.ipynb
ADDED
The diff for this file is too large to render.
See raw diff
|
|
README.md
ADDED
@@ -0,0 +1,135 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Lightweight Learned Image Compression (LLIC)
|
2 |
+
|
3 |
+
## Installation
|
4 |
+
|
5 |
+
1. Follow the installation instructions for [torch](https://pytorch.org/get-started/locally/) and [compressai](https://interdigitalinc.github.io/CompressAI/installation.html)
|
6 |
+
2. Install LLIC via pip: `pip install LLIC`
|
7 |
+
|
8 |
+
## Pre-trained checkpoints
|
9 |
+
|
10 |
+
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)
|
11 |
+
|
12 |
+
[Request access to other checkpoints (grayscale, hyperspectral, microscopy, etc)](mailto:[email protected])
|
13 |
+
|
14 |
+
## Usage example
|
15 |
+
|
16 |
+
|
17 |
+
```python
|
18 |
+
import torch
|
19 |
+
import zlib
|
20 |
+
import numpy as np
|
21 |
+
import compressai
|
22 |
+
from io import BytesIO
|
23 |
+
from IPython.display import display
|
24 |
+
from PIL import Image
|
25 |
+
from LLIC import LLIC
|
26 |
+
from torchvision.transforms import ToPILImage, PILToTensor
|
27 |
+
```
|
28 |
+
|
29 |
+
Load the model
|
30 |
+
|
31 |
+
|
32 |
+
```python
|
33 |
+
checkpoint = torch.load("LLIC_rgb_v0.0.1.pth",map_location="cpu")
|
34 |
+
codec = LLIC.RateDistortionAutoEncoder()
|
35 |
+
codec.load_state_dict(checkpoint['model_state_dict'])
|
36 |
+
```
|
37 |
+
|
38 |
+
|
39 |
+
|
40 |
+
|
41 |
+
<All keys matched successfully>
|
42 |
+
|
43 |
+
|
44 |
+
|
45 |
+
Download example image
|
46 |
+
|
47 |
+
|
48 |
+
```python
|
49 |
+
!wget https://r0k.us/graphics/kodak/kodak/kodim05.png
|
50 |
+
```
|
51 |
+
|
52 |
+
|
53 |
+
```python
|
54 |
+
original_image = Image.open("kodim05.png")
|
55 |
+
original_image
|
56 |
+
```
|
57 |
+
|
58 |
+
|
59 |
+
|
60 |
+
|
61 |
+
|
62 |
+
![png](README_files/README_6_0.png)
|
63 |
+
|
64 |
+
|
65 |
+
|
66 |
+
|
67 |
+
The analysis and synthesis transforms expect dimensions to be multiples of of 16. Zero padding can be applied otherwise.
|
68 |
+
|
69 |
+
|
70 |
+
```python
|
71 |
+
def pad(x, p=2**5):
|
72 |
+
h, w = x.size(2), x.size(3)
|
73 |
+
pad, _ = compressai.ops.compute_padding(h, w, min_div=p)
|
74 |
+
return torch.nn.functional.pad(x, pad, mode="constant", value=0)
|
75 |
+
|
76 |
+
def preprocess(pil_image):
|
77 |
+
tensor = PILToTensor()(pil_image)
|
78 |
+
tensor = tensor.unsqueeze(0)
|
79 |
+
tensor = tensor.to(torch.float)
|
80 |
+
tensor = tensor/255
|
81 |
+
tensor = tensor - 0.5
|
82 |
+
return pad(tensor)
|
83 |
+
```
|
84 |
+
|
85 |
+
Compress the image and save file
|
86 |
+
|
87 |
+
|
88 |
+
```python
|
89 |
+
padded_image = preprocess(original_image)
|
90 |
+
original_size = padded_image.shape
|
91 |
+
compressed_image, compressed_shape = LLIC.compress(padded_image, codec)
|
92 |
+
with open("kodim05.llic", 'wb') as f:
|
93 |
+
f.write(compressed_image)
|
94 |
+
```
|
95 |
+
|
96 |
+
Decompress and view the image
|
97 |
+
|
98 |
+
|
99 |
+
```python
|
100 |
+
def crop(x, size):
|
101 |
+
H, W = x.size(2), x.size(3)
|
102 |
+
h, w = size
|
103 |
+
_, unpad = compressai.ops.compute_padding(h, w, out_h=H, out_w=W)
|
104 |
+
return torch.nn.functional.pad(x, unpad, mode="constant", value=0)
|
105 |
+
|
106 |
+
def postprocess(tensor):
|
107 |
+
tensor = tensor[0] + 0.5
|
108 |
+
tensor = 255*tensor
|
109 |
+
tensor = tensor.clamp(0,255)
|
110 |
+
tensor = tensor.to(torch.uint8)
|
111 |
+
pil_image = ToPILImage()(tensor)
|
112 |
+
return pil_image
|
113 |
+
```
|
114 |
+
|
115 |
+
|
116 |
+
```python
|
117 |
+
with open("kodim05.llic", 'rb') as f:
|
118 |
+
compressed_image = f.read()
|
119 |
+
tensor = LLIC.decompress(compressed_image, compressed_shape, codec)
|
120 |
+
recovered_image = postprocess(crop(tensor, (512,768)))
|
121 |
+
```
|
122 |
+
|
123 |
+
|
124 |
+
```python
|
125 |
+
recovered_image
|
126 |
+
```
|
127 |
+
|
128 |
+
|
129 |
+
|
130 |
+
|
131 |
+
|
132 |
+
![png](README_files/README_14_0.png)
|
133 |
+
|
134 |
+
|
135 |
+
|
README_files/README_14_0.jpg
ADDED
Git LFS Details
|
README_files/README_14_0.png
ADDED
Git LFS Details
|
README_files/README_6_0.jpg
ADDED
Git LFS Details
|
README_files/README_6_0.png
ADDED
Git LFS Details
|