import matplotlib.pyplot as plt
from skimage.restoration import denoise_wavelet
from skimage.io import imread
from skimage.metrics import peak_signal_noise_ratio as compare_psnr

# 读取含有噪音的图像
# 使用 io.imread 读取图像文件 "cat_noisy.jpg"，as_gray 参数设置为 True 表示以灰度图形式读取
noisy_img = imread("img/cat_noisy.jpg", as_gray=True)
noise_Free= imread("img/cat.png", as_gray=True)

#利用for循环获取PSNR值列表并获得最大峰值信噪比的值，及其对应的sigma的值
Sugma1=np.linspace(0,1,200)
Sugma1[0]=0.005
psnr1=[]
for i in Sugma1:
    img = denoise_wavelet(noisy_img,sigma=i)
    psnr=compare_psnr(noise_Free,img,data_range=None)
    psnr1.append(psnr)
psnr_max=np.max(psnr1)
psnr_max_index=psnr1.index(psnr_max)
sigma_best=Sugma1[psnr_max_index].round(2)
psnr_max=np.round(psnr_max,2)#保留小数点后两位

# 利用小波去噪
# 使用 denoise_wavelet 方法对含有噪音的图像 noisy_img 进行小波去噪
wavelet_img = denoise_wavelet(noisy_img,sigma=sigma_best)

# 使用 matplotlib 绘图
# 创建 1 行 2 列的子图，figsize 设置图像的尺寸
fg, ax = plt.subplots(1, 4, figsize=(20, 3))

# 在第一个子图中展示原始的无噪音图像
ax[0].imshow(noise_Free, cmap='gray')
ax[0].set_title('noise_Free')  # 设置子图标题

# 在第2个子图中展示原始的噪音图像
ax[1].imshow(noisy_img, cmap='gray')
ax[1].set_title('Noisy Image')  # 设置子图标题

# 在第3个子图中展示小波去噪后的结果
ax[2].imshow(wavelet_img, cmap='gray')
ax[2].set_title(f"Wavelet Denoised (Best Sigma:{sigma_best},PSNR:{psnr_max})")  # 设置子图标题

#在第四个子图中展示PSNR与Sigma的曲线图
ax[3].plot(Sugma1,psnr1)
ax[3].set_title('PSNR VS Sigma')
ax[3].annotate(f'({sigma_best},{psnr_max})',xy=(sigma_best,psnr_max),xytext=(sigma_best+0.2,psnr_max+0.05),color="r",
               arrowprops=dict(arrowstyle="->", connectionstyle="arc3", color="r"))#在图像上对指定点进行文本注释

# 显示图像 
plt.show()





from skimage import io,img_as_float
import ipywidgets
from skimage.metrics import peak_signal_noise_ratio as compare_psnr
from skimage.restoration import (denoise_tv_chambolle, denoise_bilateral, denoise_wavelet, denoise_nl_means)
import torch
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from scipy import stats


def fft2_denoise(im):
    # TODO 完成fft2去噪算法，实现输入为原图，输出为去噪后的图像
    h, w = im.shape
    # 2. 获得平移的傅里叶系数
    fCoef = np.fft.fftshift(np.fft.fft2(im))

    # 3. 获得振幅谱
    amp_im_fft = np.log(np.abs(fCoef))

    # 4. 生成一个高斯掩模
    sigma = 0.4  # 开口大小
    new_with = max(h, w)  # 2D高斯分布必须是正方形，所以选择最长边为高斯分布标准
    xr = stats.zscore(np.arange(new_with))
    [x, y] = np.meshgrid(xr, xr)  # 获取标准正态分布2D平面上的所有点的坐标
    mask = np.exp(-(x**2 + y**2) / (2 * sigma**2))
    # 因为生成的高斯分布存在于最长边构成的正方形中，所以最后需要把多余的部分截取
    if h > w:
        wmin = (h - w) // 2
        wmax = wmin + w
        mask = mask[:, wmin:wmax]
    else:
        hmin = (w - h) // 2
        hmax = hmin + h
        mask = mask[hmin:hmax, :]

     # 5. 用掩模掩盖傅里叶系数
    fCoef_masked = fCoef * mask

    # 6. 通过逆傅里叶变换还原被“掩盖”后的傅里叶系数[傅里叶系数需要先逆平移]
    im_denoised = np.fft.ifft2(np.fft.ifftshift(fCoef_masked)).real
    return im_denoised

def optimize_denoise(img_list):
    # 以下虽然用了as_gray=True, 但是对于某些tif仍然无法转化为float
    # 所以仍然要使用img_as_float
    clean_img=img_as_float(io.imread(img_list[0],as_gray=True)) # 干净原图
    noisy_img=img_as_float(io.imread(img_list[1],as_gray=True)) # 有噪音图像
    #print(clean_img)
    # TODO 完成此部分代码实现题目要求功能
    # 2. 双边滤波去噪
    bilateral_img = denoise_bilateral(noisy_img) 
    # 3. 全变分去噪
    tv_img = denoise_tv_chambolle(noisy_img)

    # 4. 小波去噪
    wavelet_img = denoise_wavelet(noisy_img)

    # 5. Skimage的非局部均值去噪
    NLM_skimg_denoise_img = denoise_nl_means(noisy_img, patch_size=9, patch_distance=5)
    
    # 1. fft2 去噪
    fft2_img = fft2_denoise(noisy_img)

    #定义局部变量
    bilateral_psnr=0
    tv_psnr=0
    wavelet_psnr=0
    NLM_SK_denoise_psnr=0
    fft2_psnr=0
    #计算峰噪比
    bilateral_psnr = compare_psnr(clean_img, bilateral_img)
    tv_psnr = compare_psnr(clean_img, tv_img)
    wavelet_psnr = compare_psnr(clean_img, wavelet_img)
    NLM_SK_denoise_psnr = compare_psnr(clean_img, NLM_skimg_denoise_img)
    fft2_psnr = compare_psnr(clean_img, fft2_img)
    #print(bilateral_psnr)
    print(f'bilateral PSNR:{bilateral_psnr}\ntv PSNR:{tv_psnr}\nwavelet PSNR:{wavelet_psnr}\nNLM PSNR:{NLM_SK_denoise_psnr}\nfft2 PSNR:{fft2_psnr}')

    #创建数据并关联
    X = ['bilateral_img','tv_img','wavelet_img','NLM_skimg_denoise_img','fft2_img']
    Y = [bilateral_psnr,tv_psnr,wavelet_psnr,NLM_SK_denoise_psnr,fft2_psnr]
    psnr_list = pd.Series(Y,index = X)
    #print(psnr_list)
    psnr_best = psnr_list.idxmax()
    #print(img_best)
    

    fg,ax = plt.subplots(1,3,figsize=(18,6))
    ax[0].imshow(clean_img,cmap='gray')
    ax[0].set_title('clean_img')
    ax[1].imshow(noisy_img,cmap='gray')
    ax[1].set_title('noisy_img')
    
    if psnr_best == psnr_list.index[0]:
        ax[2].imshow(bilateral_img,cmap='gray')
        ax[2].set_title(f'bilateral:{psnr_list[0]}')
    if psnr_best ==psnr_list.index[1]:
        ax[2].imshow(tv_img,cmap='gray')
        ax[2].set_title(f'tv:{psnr_list[1]}')
    if psnr_best == psnr_list.index[2]:
        ax[2].imshow(wavelet_img,cmap='gray')
        ax[2].set_title(f'wavelet_img:{psnr_list[2]}')
    if psnr_best == psnr_list.index[3]:
        ax[2].imshow(NLM_skimg_denoise_img,cmap='gray')
        ax[2].set_title(f'NLM:{psnr_list[3]}')
    if psnr_best == psnr_list.index[4]:
        ax[2].imshow(fft2_img,cmap='gray')
        ax[2].set_title(f'fft2:{psnr_list[3]}')


x = ipywidgets.Dropdown(
    options=[('MRI', ("img/MRI_clean.tif","img/MRI_noisy.tif")),
             ('cat', ("img/cat.png","img/cat_noisy.jpg")), 
             ('balloon', ("img/balloon.jpg","img/balloon_noisy.jpg")), 
             ('einstein', ("img/einstein.jpg","img/einstein_noisy.jpg")),
             ]
    )
ipywidgets.interact(optimize_denoise, img_list=x)