二维图像haar小波变换的分解与重构

二维离散小波的理论推导和一维小波类似，但是以其尺度函数生成的尺度函数集作为标准正交基的尺度空间V_i的正交补空间W_i不能直接得到，而是可以证明，正交补空间W_i是由三个子空间的直和组成，对应的三个子空间可以由作为正交基的尺度函数、小波函数张成。

二维离散小波变换对图像的分解可以看做如下图所示的滤波过程，即首先进行行滤波，沿着列方向进行，然后下采样，然后对上一步得到的结果进行列滤波，沿着行方向，然后下采样，做完所有的列滤波后，获得4个不同的频带，一个近似分量、三个细节分量（水平、垂直、对角线），将所有的结果组合为一张图。

若对所得的近似分量继续进行这样的滤波过程，既可以得到如下图所示的塔式分解：

注意这里有个坑，python cv2对于float类型会乘以255，导致近似分量会显示为白图，需要使用pyplot进行显示。重构出来的图像如下图所示：

haar代码如下：

import numpy as np
import matplotlib.image as mpimg # mpimg 用于读取图片
import matplotlib.pyplot as plt

L_D = np.array([0.7071,   0.7071])
H_D = np.array([-0.7071,  0.7071])
L_R = np.array([0.7071,   0.7071])
H_R = np.array([0.7071,  -0.7071])

QP = 20
threshold = 5
#这里的图像是整张图像，height、width是整张图像大小
def haar_decomposition(img, height, width, depth):
    img_h = np.zeros((height , width ))
    img_l = np.zeros((height , width ))

    img_l_l = np.zeros((height , width)) #近似系数，左上
    img_l_h = np.zeros((height , width)) #细节系数，垂直方向，左下
    img_h_l = np.zeros((height , width)) #细节系数，水平方向，右上
    img_h_h = np.zeros((height , width)) #细节系数，对角线，右下

    img_new = np.zeros((height , width)) 

    depth_count = 1;
    while depth_count <= depth:

        if depth_count != 1:
            width = width // 2
            height = height // 2
        
        #行滤波
        #行间滤波,然后进行下采样
        for i in range(0, height // 2, 1):
            for j in range(0, width, 1):    
                i_endflag = 1 if i + 1 < height // 2 else -1
                
                img_l[i][j] = ((L_D[0] * img[i * 2][j] + L_D[1] * img[i * 2 + i_endflag][j]) // QP) * QP
                img_h[i][j] = ((H_D[0] * img[i * 2][j] + H_D[1] * img[i * 2 + i_endflag][j]) // QP) * QP
               
        n_w = width // 2
        n_h = height // 2
        #列滤波，列间滤波，沿着行
        for i in range(0, height // 2, 1):
            for j in range(0, width // 2, 1):
                j_endflag = 1 if j + 1 < width // 2 else -1
                img_l_l[i][j]             = ((L_D[0] * img_l[i][j * 2] + L_D[1] * img_l[i][j * 2 + j_endflag]) // QP) * QP
                img_l_h[n_h + i][j]       = ((H_D[0] * img_l[i][j * 2] + H_D[1] * img_l[i][j * 2 + j_endflag]) // QP) * QP
                img_h_l[i][n_w + j]       = ((L_D[0] * img_h[i][j * 2] + L_D[1] * img_h[i][j * 2 + j_endflag]) // QP) * QP
                img_h_h[n_h + i][n_w + j] = ((H_D[0] * img_h[i][j * 2] + H_D[1] * img_h[i][j * 2 + j_endflag]) // QP) * QP
    
        if depth_count == depth:
            img_new = img_new + img_l_l + img_l_h + img_h_l + img_h_h
        else:
            img_new = img_new + img_l_h + img_h_l + img_h_h
        
        img = img_l_l.copy()

        img_l[:][:] = 0
        img_h[:][:] = 0
        img_l_l[:][:] = 0    
        img_l_h[:][:] = 0
        img_h_l[:][:] = 0
        img_h_h[:][:] = 0
        
        depth_count += 1 
   
    img_new = np.round(img_new)
    
    return img_new
    
#这里的图像是整张图像，height、width是原始图像大小
def haar_reconstruction(img, height, width, trans_num):
    img_h = np.zeros((height , width ))
    img_l = np.zeros((height , width ))

    img_l_l = np.zeros((height , width))
    img_l_h = np.zeros((height , width))
    img_h_h = np.zeros((height , width))
    img_h_l = np.zeros((height , width))

    img_new = np.zeros((height , width)) 

    depth_count = 1;
    
    min_width  = width  // ( 2 ** trans_num )
    min_height = height // ( 2 ** trans_num )
    
    cur_width =  min_width
    cur_height = min_height
    
    while depth_count <= trans_num:

        if depth_count != 1:
            cur_width  = cur_width * 2
            cur_height = cur_height * 2
        
        #列上采样
        for i in range(cur_height):
            for j in range(cur_width):
                img_l_l[i][2*j] = img[i][j]
                img_l_h[i][2*j] = img[cur_height + i][j]
                img_h_l[i][2*j] = img[i][cur_width + j]
                img_h_h[i][2*j] = img[cur_height + i][cur_width + j]
        
        #列滤波,列间滤波，沿着行
        for i in range(cur_height):
            for j in range(2 * cur_width):
                j_endflag = 1 if j + 1 < 2 * cur_width else -1
                
                img_l_l[i][j] = L_R[0] * img_l_l[i][j] + L_R[1] * img_l_l[i][j + j_endflag]
                img_l_h[i][j] = H_R[0] * img_l_h[i][j] + H_R[1] * img_l_h[i][j + j_endflag]
                img_h_l[i][j] = L_R[0] * img_h_l[i][j] + L_R[1] * img_h_l[i][j + j_endflag]
                img_h_h[i][j] = H_R[0] * img_h_h[i][j] + H_R[1] * img_h_h[i][j + j_endflag]
        
        for i in range(cur_height):
            for j in range(2 * cur_width):
                img_l[i][j] = img_l_l[i][j] + img_l_h[i][j]
                img_h[i][j] = img_h_l[i][j] + img_h_h[i][j]
        
                
        #行上采样
        for i in range(cur_height - 1, -1, -1):
            for j in range(2 * cur_width):
                img_l[2 * i][j] = img_l[i][j]
                img_h[2 * i][j] = img_h[i][j]
                if i != 0:
                    img_h[i][j]     = 0
                    img_l[i][j]     = 0
                    
        #行滤波
        for j in range(2 * cur_width):
            for i in range(2 * cur_height):
                i_endflag = 1 if i + 1 < 2 * cur_height else -1
                img_l[i][j] = L_R[0] * img_l[i][j] + L_R[1] * img_l[i + i_endflag][j]
                img_h[i][j] = H_R[0] * img_h[i][j] + H_R[1] * img_h[i + i_endflag][j]

        img_new = img_l + img_h    
        img_new = np.round(img_new)
        img_new = np.clip(img_new, 0, 255)    
        
        for i in range(2 * cur_height):
            for j in range(2 * cur_width):
               img[i][j] = img_new[i][j]
        
        img_l[:][:] = 0
        img_h[:][:] = 0
        img_l_l[:][:] = 0    
        img_l_h[:][:] = 0
        img_h_l[:][:] = 0
        img_h_h[:][:] = 0
        
        depth_count += 1      
    
    return img_new

def calc_psnr(img_reco, img_org):
    data1 = np.round(img_reco)
    data2 = np.round(img_org)
    data1 = np.clip(data1, 0, 255)
    data2 = np.clip(data2, 0, 255)
    
    mse = np.mean((data1 - data2) * (data1-data2))
    # print(mse)
    return 10 * np.log10(255*255.0/mse)

# 统计零值占比
def calc_percent(img):
    count = 0
    height = img.shape[0]
    width = img.shape[1]
    for i in range(height):
        for j in range( width):
            if img[i, j] == 0:
                count += 1

    percent = count / 256 / 256
    return percent


def main():

    img = mpimg.imread("LENA.BMP")
    
    img = img.copy()
    trans_num = 1
    #img_st, 存储小波变换后的图像
    img_st = haar_decomposition(img, img.shape[0], img.shape[1], trans_num)
    plt.xticks([])
    plt.yticks([])
    plt.imshow(img_st, cmap='gray')
    plt.show()

    per = calc_percent(img_st)
    
    img_re = haar_reconstruction(img_st, img.shape[0], img.shape[1], trans_num)
    mse = calc_psnr(img, img_re)

    plt.imshow(img_re, cmap='gray')
    plt.show()
    
    print(mse)
    print(per)
    
if __name__ == "__main__":
    main()