# @File         :  PerfectFusion.py
# @Author       :  ShawnWang
# @Desc         :  多焦点图像融合
# Time          :  2023/9/24 08:25
import cv2
import matplotlib.pyplot as plt
import numpy as np
import pywt
from PIL import Image

# 基于小波变换的多聚焦图像融合算法

'''
论文中的两种方案：
1，对低频分量的所有像素点计算其局部方差，每张图所有点的方差加起来除以两张图所有点加起来，得到两张图的权重
融合图像每个像素点的值为两张图对应像素点的值加权平均，这个权就是上面算出来的权值。
2，效果不是很好，在小波分块的边缘有明显的灰度跳变(就是有些论文里说的分块效应)，but why?
3，已找到原因，要求的不是某一个点对全图的方差，而是在某点附近开个小窗口求窗口的局部方差
4，现用局部方差的方法对多聚焦图像效果非常完美
'''


# 读取图片文件 将其转化成灰度图像，同时转化成包含图像像素值的NumPy数组
def imgOpen(path):
    # Image.open(path) 会返回一个Image对象，表示打开的图像。然后，.convert('L') 方法会将该图像转换为灰度图像
    img = Image.open(path).convert('L')
    # img是一个PIL图像对象。np.array()函数用于将该图像对象转换为NumPy数组。转换后，imgArray就是一个包含图像像素值的NumPy数组。
    # 通过将图像转换为NumPy数组，可以方便地对图像进行各种数值计算和图像处理操作。可以使用NumPy提供的丰富功能来访问、修改和分析图像数据。
    imgArray = np.array(img)
    return imgArray


# 对于低频分量，计算两图的权重比
def varianceWeight(img1, img2):
    # 第一个元素mean1表示图像的均值，第二个元素var1表示图像的方差。
    mean1, var1 = cv2.meanStdDev(img1)
    mean2, var2 = cv2.meanStdDev(img2)
    weight1 = var1 / (var1 + var2)
    weight2 = var2 / (var1 + var2)
    return weight1, weight2


# 用于计算输入数组中每个像素点周围5个像素内的方差(提取一个大小为11x11的窗口)，并返回生成的方差图像
def getVarianceImg(array):
    row, col = array.shape
    varImg = np.zeros((row, col))
    for i in range(row):
        for j in range(col):
            up = i - 5 if i - 5 > 0 else 0
            down = i + 5 if i + 5 < row else row
            left = j - 5 if j - 5 > 0 else 0
            right = j + 5 if j + 5 < col else col
            # array[up:down, left:right]表示从数组的第up行（包含）到第down行（不包含），以及从第left列（包含）到第right列（不包含）的区域，即提取了一个11x11的子窗口
            window = array[up:down, left:right]
            # 利用cv2.meanStdDev()函数计算窗口中像素值的平均值和标准差。该函数会返回一个元组，其中第一个元素是窗口中像素值的平均值，第二个元素是窗口中像素值的标准差
            mean, var = cv2.meanStdDev(window)
            varImg[i, j] = var
    return varImg


def testWave(img1, img2):
    # 使用PyWavelets库中的小波变换函数对图像进行二维小波分解的操作，通过进行小波分解，可以将图像转换到频域，并提取不同频率上的信息。这对于图像处理任务如去噪、压缩和特征提取等方面非常有用
    # pywt.wavedec2() 函数以 'haar' 作为小波基函数（可根据需要更改为其他小波基函数），并指定了分解的级数为4
    # 分解的层数就越多，也就意味着能够分析的频率范围就越广。这种分解可以理解为将图像从空间域转换到频域，从而可以更好地分析图像的频率特性。。
    # 结果，transf1和transf2将是包含小波分解系数的元组(tuple)数据结构，其中每个元素都是一个二维数组（矩阵），表示对应尺度的小波系数。最顶层的元素是低频部分，而较低层的元素是高频部分。
    transf1 = pywt.wavedec2(img1, 'haar', level=4)
    transf2 = pywt.wavedec2(img2, 'haar', level=4)
    # 确保变量transf1和transf2具有相同的长度。
    assert len(transf1) == len(transf2)
    # 创建了一个空列表recWave，用于存储重构后的小波系数
    recWave = []
    for k in range(len(transf1)):
        # 处理低频分量
        if k == 0:
            # 调用varianceWeight函数，计算transf1[0] 和transf2[0]之间的权重比，并将结果分配给loWeight1和loWeight2两个变量。
            loWeight1, loWeight2 = varianceWeight(transf1[0], transf2[0])
            # 创建一个形状与transf2[0]相同的全零数组lowFreq，用于存储加权合并后的低频分量。
            lowFreq = np.zeros(transf2[0].shape)
            # 获取transf1[0]的行数和列数，并将它们分别赋值给变量row和col
            row, col = transf1[0].shape
            # 使用权重对低频分量进行加权融合，并将结果添加到recWave列表中
            for i in range(row):
                for j in range(col):
                    # 加权
                    lowFreq[i, j] = loWeight1 * transf1[0][i, j] + loWeight2 * transf2[0][i, j]
            recWave.append(lowFreq)
            continue
        # 处理高频分量
        cvtArray = []
        # 使用了zip() 函数将transf1[k] 和transf2[k] 两个列表中的元素进行逐个配对。
        # 在每一次迭代中，array1和array2分别表示来自transf1[k]和transf2[k]的配对元素
        for array1, array2 in zip(transf1[k], transf2[k]):
            tmp_row, tmp_col = array1.shape
            highFreq = np.zeros((tmp_row, tmp_col))
            # 调用getVarianceImg()函数分别计算当前数组元素array1和array2的方差图像，保存在var1和var2变量中
            var1 = getVarianceImg(array1)
            var2 = getVarianceImg(array2)
            # 双重循环遍历当前数组元素的每个像素点，根据对应位置上的方差值判断选择使用array1还是array2的像素值，并将选取的值赋给highFreq[i, j]
            for i in range(tmp_row):
                for j in range(tmp_col):
                    # 目的：在图像处理中，方差常被用作衡量像素值变化程度的指标。较大的方差通常表示像素值在该位置上的变化较大，可能对应着图像中的高频信息（如边缘、纹理等）。而较小的方差则表示像素值变化较小，对应着图像的低频信息（如平坦区域）。
                    # 通过比较var1[i, j]和var2[i, j]的大小，可以判断哪个数组元素所对应位置的像素值变化较大，即具有更高的频率成分。
                    highFreq[i, j] = array1[i, j] if var1[i, j] > var2[i, j] else array2[i, j]
            cvtArray.append(highFreq)
        recWave.append(tuple(cvtArray))
    return pywt.waverec2(recWave, 'haar')


# 在一个窗口中显示融合图像
def testPlot(org1, org2, img):
    # 函数中使用的子图编号（131、132、133）是一种常见的约定，用于表示一个1行3列的子图网格中的三个子图位置。
    plt.subplot(131)
    # 每个子图都使用plt.imshow()函数来显示对应的图像，并可以通过cmap='gray'参数将图像以灰度的形式进行显示。
    plt.imshow(org1, cmap='gray')
    # 调用plt.axis('off')可以去除图像周围的坐标轴标签和刻度。
    plt.axis('off')

    plt.subplot(132)
    plt.imshow(org2, cmap='gray')
    plt.axis('off')

    plt.subplot(133)
    plt.imshow(img, cmap='gray')
    plt.axis('off')

    plt.show()


if __name__ == '__main__':
    img1 = imgOpen('data/lytro-03-A.jpg')
    img2 = imgOpen('data/lytro-03-B.jpg')
    #
    # img1 = imgOpen('data/lytro-02-A.jpg')
    # img2 = imgOpen('data/lytro-02-B.jpg')

    rec = testWave(img1, img2)
    testPlot(img1, img2, rec)