# -*- coding: utf-8 -*-
"""Débruitage d’images par seuillage des coefficients de Fourier"""
import numpy as np
from matplotlib import pyplot as plt

def thresholdimage_fourier(y,T,mode):
    """Applique un seuillage sur les coefficients de Fourier d'une image
        Paramètres
        ----------
        y (matrice) : Une image
        T (flottant) : Un seuil
        mode (string) : Le mode de seuillage ('soft' ou 'hard')
    """
    # Calcul des coefficients de Fourier de l'image
    coeffs = np.fft.fft2(y)

    # Seuillage des coefficients
    if mode == 'soft':
        newcoeffs = soft_thresholding(coeffs, T)
    else:
        newcoeffs = hard_thresholding(coeffs, T)

    # Reconstruction de l'image
    NewImage = np.real(np.fft.ifft2(newcoeffs))

    # Retourne l'image après débruitage
    return NewImage

def soft_thresholding(x, T):
    """Applique un seuillage doux sur un signal
        Paramètres
        ----------
        x (vecteur): Un signal
        T (flottant) : Un seuil
    """
    signal = x.copy()
    signal[(np.abs(x) - T) < 0] = 0
    signal[(np.abs(x) - T) >= 0] -= T * np.sign(signal[(np.abs(signal) - T) >= 0])
    return signal

def hard_thresholding(x, T):
    """Applique un seuillage dur sur un signal
        Paramètres
        ----------
        x (vecteur) : Un signal
        T (flottant): Un seuil
    """
    signal = x.copy()
    signal[(np.abs(signal) - T) < 0] = 0
    return signal

def addwhitenoise(signal,sigma):
    """Ajoute un bruit blanc gaussien centré à un signal
        Paramètres
        ----------
        signal (vecteur) : Un signal
        sigma (flottant) : L'écart-type du bruit gaussien
    """
    y = signal + np.random.normal(0, sigma, size=np.shape(signal))
    return y 

# Fonction du calcul du SNR (Signal to Noise Ratio)
def snr(xtilde,x):
    """Calcule l'indice SNR de deux signaux
        ----------
        x (vecteur) : Signal de départ
        xtilde (vecteur) : Estimation d'un signal
    """
    s = 10 * np.log10(np.mean(x ** 2) / np.mean((x - xtilde) ** 2))
    return s




#Estimation de T par maximisation du SNR

# On charge ici une image en noire et blanc de taille (512, 512)
img = plt.imread("C:/Users/Administrator/Desktop/漆膜数据集样本/Image-Denoising-master/barb.bmp")
# On convertit l'image en flottant pour pouvoir effectuer nos calculs
img = img.astype(float)
img = img / 255

# Bruitage de l'image
sig = 0.1
y = addwhitenoise(img, sig)
T = sig

# Seuils et SNR
T = np.linspace(0, 200, 500)
SNR_soft = np.zeros(len(T))
SNR_hard = np.zeros(len(T))

# Calcule des valeurs du SNR en fonction de T (Seuillage doux)
for i in range(len(T)):
    xest = thresholdimage_fourier(y ,T = T[i] , mode = 'soft')
    SNR_soft[i] = snr(xest,img)

# Calcule des valeurs du SNR en fonction de T (Seuillage dur)
for i in range(len(T)):
    xest = thresholdimage_fourier(y ,T = T[i] , mode = 'hard')
    SNR_hard[i] = snr(xest,img)

# Estimation de la valeur de T pour laquelle le SNR est maximal
Tmax_soft = np.around(T[np.argmax(SNR_soft)],2)
Tmax_hard = np.around(T[np.argmax(SNR_hard)],2)

# Estimation de la valeur de T/sig pour laquelle le SNR est maximal
T_sig_max_soft = np.around(Tmax_soft / sig,2)
T_sig_max_hard = np.around(Tmax_hard / sig,2)

print('(Seuillage doux) Le SNR est maximal pour T =',Tmax_soft,2)
print('(Seuillage doux) Le SNR est maximal pour T/sigma =',T_sig_max_soft)
print('(Seuillage doux) SNRmax = ',np.around(np.amax(SNR_soft),2))
print('----------')
print('(Seuillage dur) Le SNR est maximal pour T =',Tmax_hard)
print('(Seuillage dur) Le SNR est maximal pour T/sigma =',T_sig_max_hard)
print('(Seuillage dur) SNRmax = ',np.around(np.amax(SNR_hard),2))


# Affichage des images
fig, axarr = plt.subplots(3, 2, figsize=(15,10))

# Image originale
axarr[0,0].set_title('Image originale')
axarr[0,0].imshow(img, cmap="gray")

# Image bruitée
axarr[0,1].set_title('Image bruitée ($\sigma = 0.1$)')
axarr[0,1].imshow(y, cmap="gray")
axarr[0,1].text(540, 490, "SNR = " + str(np.around(snr(img,y),2)),
                bbox={'facecolor': 'white', 'pad': 5})

# Courbe du SNR en fonction de T (Seuillage doux)
axarr[1,0].set_title(r'Courbe du SNR en fonction de T (Seuillage doux)')
axarr[1,0].plot(T,SNR_soft,'k', linewidth=1, label='densité')

axarr[1,0].set_xlabel("T")
axarr[1,0].set_ylabel("SNR")
axarr[1,0].grid(True)

# Restauration par seuillage doux des coefficients de Fourier
axarr[1,1].set_title('Restauration par seuillage doux des coefficients de Fourier (T = ' + str(Tmax_soft) + ')')
axarr[1,1].imshow(thresholdimage_fourier(y,T = Tmax_soft ,mode = 'soft'), cmap="gray")
axarr[1,1].text(540, 490, "SNR = " + str(np.around(np.amax(SNR_soft),2)),
                bbox={'facecolor': 'white', 'pad': 5})

# Courbe du SNR en fonction de T (Seuillage dur)
axarr[2,0].set_title(r'Courbe du SNR en fonction de T (Seuillage dur)')
axarr[2,0].plot(T,SNR_hard,'k', linewidth=1, label='densité')

axarr[2,0].set_xlabel("T")
axarr[2,0].set_ylabel("SNR")
axarr[2,0].grid(True)

# Restauration par seuillage dur des coefficients de Fourier
axarr[2,1].set_title('Restauration par seuillage dur des coefficients de Fourier (T = ' + str(Tmax_hard) + ')')
axarr[2,1].imshow(thresholdimage_fourier(y,T= Tmax_hard ,mode = 'hard'), cmap="gray")
axarr[2,1].text(540, 490, "SNR = " + str(np.around(np.amax(SNR_hard),2)),
                bbox={'facecolor': 'white', 'pad': 5})
plt.subplots_adjust(top = 1.5)
plt.show()