MS-SSIM图像评估算法Paddle版本复现(完善中)

前置工作

MS-SSIM需要用到SSIM，为SSIM的改进版本
python numpy版本SSIM复现与算法解释

算法公式

算法实现

ssim使用paddle官方使用高斯滤波的版本进行修改

def _ssim(img1, img2, L = 255):
    """Calculate SSIM (structural similarity) for one channel images.
    It is called by func:`calculate_ssim`.
    Args:
        img1 (ndarray): Images with range [0, 255] with order 'HWC'.
        img2 (ndarray): Images with range [0, 255] with order 'HWC'.
    Returns:
        float: ssim result.
    """
    C1 = (0.01 * L)**2
    C2 = (0.03 * L)**2

    img1 = img1.astype(np.float64)
    img2 = img2.astype(np.float64)
    kernel = cv2.getGaussianKernel(11, 1.5)
    window = np.outer(kernel, kernel.transpose())

    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]
    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
    mu1_sq = mu1**2
    mu2_sq = mu2**2
    mu1_mu2 = mu1 * mu2
    sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
    sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2

    ssim_map = ((2 * mu1_mu2 + C1) *
                (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
                                       (sigma1_sq + sigma2_sq + C2))
    cs_map = np.divide((2*sigma12+C2), (sigma1_sq+sigma2_sq+C2))
    return ssim_map.mean(), cs_map.mean()

def mul_ssim(img1, img2):
    if len(img1.shape) > 2:
        ssim_arr = []
        cs_arr = []
        for i in range(0, img1.shape[2]):
            ssim_val, cs = _ssim(img1[:, :, i], img2[:, :, i])
            ssim_arr.append(ssim_val)
            cs_arr.append(cs)
        return np.array(ssim_arr).mean(), np.array(cs_arr).mean()
    else:
        return _ssim(img1, img2)

def _msssim(img1, img2):
    weight = [0.0448, 0.2856, 0.3001, 0.2363, 0.1333]
    level = len(weight)
    mcs_array = []
    ssim_array = []
    img1 = np.array([[img1]])
    img2 = np.array([[img2]])
    for i in range(level):
        ssim, cs = mul_ssim(img1[0][0], img2[0][0])
        mcs_array.append(cs.mean())
        ssim_array.append(ssim)
        with fluid.dygraph.guard():
            padding = (img1.shape[2]%2, img2.shape[3]%2)
            pool2d_1 = fluid.dygraph.Pool2D(pool_type='avg', pool_size=2, use_cudnn=False, pool_padding=padding)
            pool2d_2 = fluid.dygraph.Pool2D(pool_type='avg', pool_size=2, use_cudnn=False, pool_padding=padding)
            filtered_im1 = pool2d_1(to_variable(img1))
            filtered_im2 = pool2d_2(to_variable(img2))
            img1 = filtered_im1.numpy()
            img2 = filtered_im2.numpy()

    overall_mssim = np.prod(np.power(mcs_array[:level - 1], weight[:level - 1])) * (
                ssim_array[level - 1] ** weight[level - 1])
    return overall_mssim