Python 二维离散傅里叶变换

Python 二维离散傅里叶变换

需要的库

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

计算两张图片的PSNR

def PSNR(A, B):
    MSE = np.sum((A - B) ** 2) / A.shape[0] / A.shape[1]
    return 10 * np.log10((255 ** 2) / MSE)

---

二维离散傅里叶变换

def DFT_2(f, N, M):
    '''
    :param f: 2_dim marix
    :param N: even; denote number of rows
    :param M: even; denote number of cols
    :return: DFT results without shifting(complex number)
    '''
    if N % 2 != 0 or M % 2 != 0:
        print("param is wrong")
        return
    f = f.astype('float64')
    rows = f.shape[0]
    cols = f.shape[1]
    n = np.arange(0, N, 1).reshape((N, 1))
    row = np.arange(0, rows, 1).reshape((1, rows))
    left = np.exp(-1j*2*np.pi/N*(n @ row))
    col = np.arange(0, cols, 1).reshape((cols, 1))
    m = np.arange(0, M, 1).reshape((1, M))
    right = np.exp(-1j*2*np.pi/M*(col @ m))
    F = left @ f @ right
    return F

---

二维离散傅里叶逆变换

def IDFT_2(F, rows, cols, option=True):
    '''
    :param F: 2-dim freqency spectrum without shifting
    :param rows: rows of original picture
    :param cols: cols of original picture
    :param option: whether make some process for results
    :return: picture in space domain
    '''
    if cols % 2 != 0 or rows % 2 != 0:
        print("IDFT params are wrong")
        return
    N = F.shape[0]
    M = F.shape[1]
    row = np.arange(0, rows, 1).reshape((rows, 1))
    n = np.arange(0, N, 1).reshape((1, N))
    left = np.exp(1j*2*np.pi/N*(row @ n))
    m = np.arange(0, M, 1).reshape((M, 1))
    col = np.arange(0, cols, 1).reshape((1, cols))
    right = np.exp(1j*2*np.pi/M*(m @ col))
    f = (left @ F @ right) / rows / cols
    if option:
        f = np.abs(f)
        f = np.around((f - np.min(f)) / (np.max(f) - np.min(f)) * 255)
    return f

其中option为True时，返回值为复数的模，否则返回值为复数。

---

频域平移

DFT的结果具有共轭对称性，因此需要将结果平移以后才能看出结果的这种对称性。

def DFT_shifting(F):
    F_shift = 1j + np.zeros(F.shape)
    N = F.shape[0]
    M = F.shape[1]
    F_shift[N>>1:N, M>>1:M] = F[0:N>>1, 0:M>>1]
    F_shift[N>>1:N, 0:M>>1] = F[0:N>>1, M>>1:M]
    F_shift[0:N>>1, M>>1:M] = F[N>>1:N, 0:M>>1]
    F_shift[0:N>>1, 0:M>>1] = F[N>>1:N, M>>1:M]
    return F_shift

其中F表示上述DFT_2()求出的频域结果。

---

绘制频域图像

由于DFT得到的结果是复数，需要对其求模值，同时取log

def DFT_2_pic(f, N, M):
    '''
    :param f: 2_dim marix
    :param N: even; denote number of rows
    :param M: even; denote number of cols
    :return: DFT results with shifting(complex number) and make log
    '''
    if N % 2 != 0 or M % 2 != 0:
        print("param is wrong")
        return
    F = DFT_2(f, N, M)
    shift = DFT_shifting(F)
    F_pic = np.log2(np.abs(shift))
    F_pic = np.round((F_pic-np.min(F_pic)) / (np.max(F_pic) - np.min(F_pic)) * 255)   # make pixel value in 0 - 255
    F_pic = F_pic.astype('uint8')
    return shift, F_pic

其中F_pic返回取对数以后的频谱；shift表示进行平移以后的复数频谱。

---