python 实现快速傅里叶变换（二维）

二维傅里叶变换可以分解成 两个维度上分别做一次一维傅里叶变换

以下代码清楚地展示了这个性质

• DFT 是对一维向量做傅里叶变换
• FFT 是对二维矩阵的每一行做一次傅里叶变换
• FFT2D 则是调用两次 FFT，在两个维度分别做傅里叶变换

def DFT(sig):
    N = sig.size
    V = np.array([[np.exp(-1j*2*np.pi*v*y/N) for v in range(N)] for y in range(N)])
    return sig.dot(V)

def FFT(x):
    N = x.shape[1] # 只需考虑第二个维度，然后在第一个维度循环
    if N % 2 > 0:
        raise ValueError("size of x must be a power of 2")
    elif N <= 8:  # this cutoff should be optimized
        return np.array([DFT(x[i,:]) for i in range(x.shape[0])])
    else:
        X_even = FFT(x[:,::2])
        X_odd = FFT(x[:,1::2])
        factor = np.array([np.exp(-2j * np.pi * np.arange(N) / N) for i in range(x.shape[0])])
        return np.hstack([X_even + np.multiply(factor[:,:int(N/2)],X_odd),
                               X_even + np.multiply(factor[:,int(N/2):],X_odd)])

def FFT2D(img):
    return FFT(FFT(img).T).T

def FFT_SHIFT(img):
    M,N = img.shape
    M = int(M/2)
    N = int(N/2)
    return np.vstack((np.hstack((img[M:,N:],img[M:,:N])),np.hstack((img[:M,N:],img[:M,:N]))))

测试

import cv2
from numpy.fft import *

img = cv2.imread("Lenna.png", cv2.IMREAD_COLOR)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

my_fft = abs(FFT_SHIFT(FFT2D(gray)))
target = abs(fftshift(fft2(gray)))
print('distance from numpy.fft: ',np.linalg.norm(my_fft-target))

plt.subplot(2,2,1)
plt.title('Lenna')
plt.imshow(gray)
plt.subplot(2,2,2)
plt.title('my FFT2D')
plt.imshow(np.log(1+my_fft))
plt.subplot(2,2,3)
plt.title('numpy.fft2')
plt.imshow(np.log(1+target))
plt.show()

变换后的频谱图和 numpy 调用的 fft2 几乎相同，说明 FFT2D 和 FFT_SHIFT 都没有问题：

distance from numpy.fft:  2.330463500403555e-08