File size: 2,336 Bytes
83630b2
 
 
 
 
 
 
 
 
 
 
6a8d727
83630b2
 
6a8d727
83630b2
 
 
 
 
 
 
 
 
 
6a8d727
83630b2
 
 
 
 
 
6a8d727
83630b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6a8d727
83630b2
 
 
 
 
 
 
 
 
 
6a8d727
 
 
83630b2
 
 
 
 
6a8d727
83630b2
 
 
 
 
 
 
 
 
 
6a8d727
 
 
 
 
 
 
 
 
 
 
83630b2
 
 
 
 
 
6a8d727
83630b2
 
 
 
 
6a8d727
 
 
 
83630b2
6a8d727
83630b2
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
137
138
139
140
141
142
143
144
145
```python
!pip install walloc PyWavelets pytorch-wavelets
```


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


```python
import os
import torch
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
from IPython.display import display
from torchvision.transforms import ToPILImage, PILToTensor
from walloc.walloc import Walloc
class Args: pass
```


```python
device = "cpu"
checkpoint = torch.load("v0.6.1.pth",map_location="cpu")
args = checkpoint['args']
codec = Walloc(
    channels = args.channels,
    J = args.J,
    N = args.N,
    latent_dim = args.latent_dim,
    latent_bits = 5
)
codec.load_state_dict(checkpoint['model_state_dict'])
codec = codec.to(device)
```


```python
img = Image.open("kodim05.png")
img
```




    
![png](README_files/README_4_0.png)
    




```python
with torch.no_grad():
    codec.eval()
    x = PILToTensor()(img).to(torch.float)
    x = (x/255 - 0.5).unsqueeze(0).to(device)
    x_hat, _, _ = codec(x)
ToPILImage()(x_hat[0]+0.5)
```




    
![png](README_files/README_5_0.png)
    




```python
with torch.no_grad():
    codec.eval()
    X = codec.wavelet_analysis(x,J=codec.J)
    Y = codec.encoder(X)
    X_hat = codec.decoder(Y)
    x_hat = codec.wavelet_synthesis(X_hat,J=codec.J)

print(f"dimensionality reduction: {x.numel()/Y.numel()}x")
```

    dimensionality reduction: 12.0x



```python
Y.unique()
```




    tensor([-15., -14., -13., -12., -11., -10.,  -9.,  -8.,  -7.,  -6.,  -5.,  -4.,
             -3.,  -2.,  -1.,   0.,   1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,
              9.,  10.,  11.,  12.,  13.,  14.,  15.])




```python
plt.hist(Y.flatten().numpy(),range=(-17.5,17.5),bins=35);
```


    
![png](README_files/README_8_0.png)
    



```python
grid_size = 4
n_channels, H, W = Y[0].shape
combined_image = Image.new('L', (W * grid_size, H * grid_size))
size_bytes = 0
for i, channel in enumerate(Y[0]):
    channel = (channel+16).to(torch.uint8)
    row = i // grid_size
    col = i % grid_size
    channel = ToPILImage()(channel)
    combined_image.paste(channel, (col * W, row * H))
combined_image
```




    
![png](README_files/README_9_0.png)
    




```python
combined_image.save('tmp.png')
print("compression_ratio: ", x.numel()/os.path.getsize("tmp.png"))
```

    compression_ratio:  20.75383532723434