pywt

1
import numpy as np
2
from matplotlib import pyplot as plt
3

4
import pywt
5
from pywt._doc_utils import draw_2d_wp_basis, wavedec2_keys
6

7
x = pywt.data.camera().astype(np.float32)
8
shape = x.shape
9

10
max_lev = 3       # how many levels of decomposition to draw
11
label_levels = 3  # how many levels to explicitly label on the plots
12

13
fig, axes = plt.subplots(2, 4, figsize=[14, 8])
14
for level in range(max_lev + 1):
15
    if level == 0:
16
        # show the original image before decomposition
17
        axes[0, 0].set_axis_off()
18
        axes[1, 0].imshow(x, cmap=plt.cm.gray)
19
        axes[1, 0].set_title('Image')
20
        axes[1, 0].set_axis_off()
21
        continue
22

23
    # plot subband boundaries of a standard DWT basis
24
    draw_2d_wp_basis(shape, wavedec2_keys(level), ax=axes[0, level],
25
                     label_levels=label_levels)
26
    axes[0, level].set_title(f'{level} level\ndecomposition')
27

28
    # compute the 2D DWT
29
    c = pywt.wavedec2(x, 'db2', mode='periodization', level=level)
30
    # normalize each coefficient array independently for better visibility
31
    c[0] /= np.abs(c[0]).max()
32
    for detail_level in range(level):
33
        c[detail_level + 1] = [d/np.abs(d).max() for d in c[detail_level + 1]]
34
    # show the normalized coefficients
35
    arr, slices = pywt.coeffs_to_array(c)
36
    axes[1, level].imshow(arr, cmap=plt.cm.gray)
37
    axes[1, level].set_title(f'Coefficients\n({level} level)')
38
    axes[1, level].set_axis_off()
39

40
plt.tight_layout()
41
plt.show()
42