Is the high frequency component more robust to noise? (PSNR Ver.)

####################################################################################
SNR_length = 30
filtering_size = 30
####################################################################################

import acoustics

def add_AWGN_to_img(img, SNRdB) :
    
    height, width = img.shape

    noise = acoustics.generator.white(img.size).reshape(img.shape)
    SNR = 10.0**(SNRdB/10.0)
    current_SNR = np.mean(img) / np.std(noise)
    noise *= (current_SNR / SNR)

    return img + noise

def PSNR(ori_img, con_img):

  max_pixel = 255.0
  mse = np.mean((ori_img - con_img)**2)
  psnr = 20* math.log10(max_pixel / math.sqrt(mse))
  
  return round(psnr, 2)

def MSE(ori_img, con_img) :
    ori_list = ori_img.flatten().tolist()
    con_list = con_img.flatten().tolist()

    save_add = 0

    for i in range(len(ori_list)) :
        save_add += (ori_list[i] - con_list[i])**2

    return round(save_add/len(ori_img), 3)
        


def cal_D(c_row, c_col, r, c) :
    s = (c_row-r)**2+ (c_col-c)**2
    return s**(1/2)


def filter_radius(fshift, rad, low = True) :
    rows, cols = fshift.shape
    c_row, c_col = int(rows/2), int(cols/2)

    filter_fshift = fshift.copy()

    for r in range(rows) :
        for c in range(cols) :
            if low :
                if cal_D(c_row, c_col, r, c) > rad :
                    filter_fshift[r, c] = 0

            else :
                if cal_D(c_row, c_col, r, c) < rad :
                    filter_fshift[r, c] = 0

    return filter_fshift

####################################################################################

img = cv2.imread("/content/lena.png", cv2.IMREAD_GRAYSCALE)
f = np.fft.fft2(img)
fshift = np.fft.fftshift(f)

####################################################################################

low_fshift = filter_radius(fshift, rad = filtering_size, low = True)
low_ishift = np.fft.ifftshift(low_fshift)
low_img = np.abs(np.fft.ifft2(low_ishift))

####################################################################################

high_fshift = filter_radius(fshift, rad = filtering_size, low = False)
high_ishift = np.fft.ifftshift(high_fshift)
high_img = np.abs(np.fft.ifft2(high_ishift))

####################################################################################



SNRdB = []

for i in range(SNR_length) :
    SNRdB.append(i+1)

print("SNRdB = {}".format(SNRdB))

psnr_list = []
for i in range(len(SNRdB)) :
    output = add_AWGN_to_img(img, i)
    psnr_list.append(PSNR(img, output))

psnr_list_low = []
for i in range(len(SNRdB)) :
    output = add_AWGN_to_img(low_img, i)
    psnr_list_low.append(PSNR(low_img, output))

psnr_list_high = []
for i in range(len(SNRdB)) :
    output = add_AWGN_to_img(high_img, i)
    psnr_list_high.append(PSNR(high_img, output))

plt.plot(SNRdB, psnr_list, marker='o', linestyle='dashed', color='blue', label='Original Image')
plt.plot(SNRdB, psnr_list_low, marker='x', linestyle='dashed', color='red', label='LPF')
plt.plot(SNRdB, psnr_list_high, marker='o', linestyle='dashed', color='green', label='HPF')
plt.axis([0, SNRdB[-1], 0, 100])
plt.grid(True)
plt.legend()
plt.show()

plt.subplot(141),plt.imshow(img, cmap = 'gray')
plt.title('Original'), plt.xticks([]), plt.yticks([])

plt.subplot(142),plt.imshow(low_img, cmap = 'gray')
plt.title('LPF'), plt.xticks([]), plt.yticks([])

plt.subplot(143),plt.imshow(high_img, cmap = 'gray')
plt.title('HPF'), plt.xticks([]), plt.yticks([])

plt.subplot(144),plt.imshow(high_img + low_img, cmap = 'gray')
plt.title('LPF + HPF'), plt.xticks([]), plt.yticks([])

plt.show()