import numpy as np
import cv2 as cv
import matplotlib.pyplot as plt


def frequency_filter(image, filter):
    """
    :param image:
    :param filter: 频域变换函数
    :return:
    """
    fftImg = np.fft.fft2(image)  # 对图像进行傅里叶变换
    fftImgShift = np.fft.fftshift(fftImg)  # 傅里叶变换后坐标移动到图像中心
    handle_fftImgShift1 = fftImgShift*filter  # 对傅里叶变换后的图像进行频域变换

    handle_fftImgShift2 = np.fft.ifftshift(handle_fftImgShift1)
    handle_fftImgShift3 = np.fft.ifft2(handle_fftImgShift2)
    handle_fftImgShift4 = np.real(handle_fftImgShift3)  # 傅里叶反变换后取频域
    return np.uint8(handle_fftImgShift4)


# 理想低通滤波器
def ILPF(image, d0, n):
    H = np.empty_like(image, dtype=float)
    M, N = image.shape
    mid_x = int(M/2)
    mid_y = int(N/2)
    for y in range(0, M):
        for x in range(0, N):
            d = np.sqrt((x - mid_x) ** 2 + (y - mid_y) ** 2)
            if d <= d0:
                H[y, x] = 1**n
            else:
                H[y, x] = 0**n
    return H


# 巴特沃斯低通滤波器
def BLPF(image, d0, n):
    H = np.empty_like(image, float)
    M, N = image.shape
    mid_x = int(M/2)
    mid_y = int(N/2)
    for y in range(0, M):
        for x in range(0, N):
            d = np.sqrt((x - mid_x) ** 2 + (y - mid_y) ** 2)
            H[y, x] = 1/(1+(d/d0)**(n))
    return H


# 高斯低通滤波器
def GLPF(image, d0, n):
    H = np.empty_like(image, float)
    M, N = image.shape
    mid_x = M/2
    mid_y = N/2
    for x in range(0, M):
        for y in range(0, N):
            d = np.sqrt((x - mid_x)**2 + (y - mid_y) ** 2)
            H[x, y] = np.exp(-d**n/(2*d0**n))
    return H


# 读取图像
image = cv.imread('D:/Python/FrequencyDomainProcessing/img/moon.jpg')
img_RGB = cv.cvtColor(image, cv.COLOR_BGR2RGB)  # 转成RGB 方便后面显示
grayImage = cv.cvtColor(img_RGB, cv.COLOR_BGR2GRAY)

result2 = frequency_filter(grayImage, ILPF(grayImage, 60, n=1))
result3 = frequency_filter(grayImage, BLPF(grayImage, 60, n=2))
result4 = frequency_filter(grayImage, BLPF(grayImage, 90, n=2))
result5 = frequency_filter(grayImage, GLPF(grayImage, 60, n=2))
result6 = frequency_filter(grayImage, GLPF(grayImage, 90, n=2))

# 显示图形
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.subplot(321), plt.imshow(grayImage, cmap=plt.cm.gray), plt.title('原始图像'), plt.axis('off')  # 坐标轴关闭
plt.subplot(322), plt.imshow(result2, cmap=plt.cm.gray), plt.title('理想低通滤波(D0=60)'), plt.axis('off')
plt.subplot(323), plt.imshow(result3, cmap=plt.cm.gray), plt.title('Butterwoth低通(D0=60,n=2)'), plt.axis('off')
plt.subplot(324), plt.imshow(result4, cmap=plt.cm.gray), plt.title('Butterwoth低通(D0=90,n=2)'), plt.axis('off')
plt.subplot(325), plt.imshow(result5, cmap=plt.cm.gray), plt.title('Gauss低通(D0=60,n=2)'), plt.axis('off')
plt.subplot(326), plt.imshow(result6, cmap=plt.cm.gray), plt.title('Gauss低通(D0=90,n=2)'), plt.axis('off')
plt.show()