python手写实现2维FFT与DCT

题外话
好久没更新了，这段时间实在是摸鱼，人快闲出毛病来了。这是一个选修课的大作业的一部分，2D-FFT的思路是借鉴了一个博客的，但做了少许改进。DCT是自己写的，都不难。这门课咋只得了81分，破防了。
一、2维FFT简述
就不放公式了，2维FFT就是两次一维FFT。一个2维信号可以看作一个矩阵，先行再列或者先列再行都可以。示意图如下图所示：

如前所述，2维FFT编写的关键仍在1维FFT。按照蝶形流图的方式编程是比较直观的，采用按时间抽取算法，第一层是2点FFT，第二层是4点FFT，以此类推。具体的细节之处在代码注释里体现，在此就不再赘述。
值得注意的是：蝶形运算是一种规则的运算，其只适用于信号长度为2的整数次幂的信号。对于任意长度信号，理论上应当补0到2的整数次幂点进行计算。但如果利用python进行编写，则可以利用其try-except结构来简化问题。具体的细节在代码注释中有所体现。
总而言之，编程的思路为将2-d FFT分解为两次1-d FFT。在计算一维序列的FFT时，首先将索引值进行反转，然后利用蝶形运算的结构逐层运算。

代码如下：

#In[]
#FFT2

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from math import log2
import copy
from numpy.core.arrayprint import dtype_is_implied

from numpy.core.getlimits import _register_type
from numpy.lib.type_check import imag


#RGB图像转为灰度图像

def RGB_to_Grey(rgb_image):
    '''RGB转灰度图像约定俗成的公式:
    G = 0.299 R + 0.587 G + 0.114 B
    '''
    Grey = np.dot(rgb_image[...,:3],[0.299,0.587,0.114])
    return Grey



#将索引倒序,是DIT-FFT
#思想是将索引转为二进制字符,再倒序,再转为整数值即可
def index_inv(index,n):
    index_bin = bin(index)[2:]  #转为二进制,字符串类型,开头是0b因此舍去
    if len(index_bin) != n:#如果位数不相等，则需要补齐0
        index_bin = '0' * (n - len(index_bin)) + index_bin
    
    index_bin_inv = index_bin[::-1] #倒序
    return int(index_bin_inv,2)


#蝶形运算 核心部分 一维FFT
def FFT1(x_n):
    '''
    按照DIT－FFT蝶形运算的流程。N点的序列需要log2(N)层蝶形运算。
    第一层蝶形运算的元素排列是原序数对应二进制数的倒序，例如8点序列做DIT-FFT时的初始排序应为0 4 2 6 1 5 3 7
    因此第一步要做的是按照逆序数来构建新序列，然后第一层是2点FFT,第二层是4,第三层为8，以此类推
    如果N不满足N=2^k,则将原序列补0为2^k,k=int(log(N)).
    '''
    x_n = x_n.tolist()
    #print(type(x_n))
    length = len(x_n)
    N = int(log2(length))

    X_k = np.zeros((length), dtype=complex)  #创建X_k

    #首先按照逆序数进行重排序   
    for i in range(length):
        try:
            X_k[i] = x_n[index_inv(i,N)] 
        except:
            X_k[i] = 0  #点数不是2的整数次幂，直接赋0，因为补0后本来就是应该为0

    for i in range(1,N+1): #一共进行N次蝶形运算
        temp = copy.copy(X_k)  #复制一个temp,且每次蝶形运算完毕后更新
        for j in range(length): #每次蝶形运算要遍历length次
            flag = j % (pow(2,i))  #第一层是2点FFT,第二层是4,第三层为8，以此类推
            if flag < pow(2,i - 1):  #只计算一个半边
                try:
                    X_k[j] = temp[j] + temp[j + pow(2,i - 1)] * np.exp(complex(0, -2*np.pi*flag/pow(2,i))) #补零为pow(2,i)点
                    X_k[j + pow(2,i - 1)] = temp[j] - temp[j + pow(2,i - 1)] * np.exp(complex(0, -2*np.pi*flag/(pow(2,i))))
                except:  #点数不是2的整数次幂，直接过。因为会报index error，而补零后无非是多了几个点，如果做等点数的变换，这些点是不需要的。
                    #print(i,j)
                    #print(flag)
                    pass


    return X_k
#In[]
#主函数
def main():

    image = mpimg.imread('image3.jpg') #读图像
    Grey = RGB_to_Grey(image) #转为灰度值

    Size = list(Grey.shape)
    Rows,Cols = Size[0],Size[1]
    #分别对行列作FFT
    #可先逐行作FFT,再逐列FFT

    F_Image = np.zeros(Size,dtype=complex)  #创建傅里叶变换后的F_Image

    for row in range(Rows):
        F_Image[row,:] = FFT1(Grey[row,:])
    
    for col in range(Cols):
        F_Image[:,col] = FFT1(F_Image[:,col])

    F_Image_M = np.abs(F_Image) #幅度
    F_Image_M_log = np.log2(F_Image_M)

    F_IMAGE_P = np.angle(F_Image) #相位

    print(F_Image_M[3,3])

    plt.figure(1)
    plt.imshow(Grey,'gray')
    plt.title('Original image')

    plt.figure(2)
    plt.imshow(F_Image_M_log,'gray')
    plt.title('FFT spectrum')
    
    plt.figure(3)
    plt.imshow(F_IMAGE_P,'gray')
    plt.title('FFT phase')




if __name__ == '__main__':
    main()          

  \space  
仿真结果：
原图片：

幅值：

相位：

二、2维DCT简述
DCT，即离散余弦变换。由于傅里叶变换的参数都是复数，因此其占用的存储空间是实数的两倍，且计算复杂。为解决这个问题，引入了离散余弦变换。也可以认为离散余弦变换是离散傅里叶变换的实数部分（是构造了实偶函数取实部所得）。
二维DCT的表达式如下：

大多数情形下，

因此，二维DCT的实现与二维FFT(DFT)相同，都可以分解为两次一维变换。
更方便的表示形式是用矩阵表示。一维DCT的系数矩阵为：

观察二维DCT变换的表达式，容易得出，变换后矩阵
$$T=Af{A}^T$$
其中$T=[F(u,v){]}_{M\times N},f=[f(x,y){]}_{M\times N}$,这可以作为DCT编程的基本思想。

代码如下：

#In[]
#DCT
#尝试用矩阵方式来作DCT变换
'''
    对于MxN的图像,假设灰度矩阵为f,shape为MxN,则DCT变换可表示为
    T = AfB,其中A:MxM,B:NxN.这样相当于先对f的列作1d-DCT,再对行作1d-DCT.
    关键是找到A和B
'''

import numpy as np
from math import sqrt,cos
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

def RGB_to_Grey(rgb_image):
    '''RGB转灰度图像约定俗成的公式:
    G = 0.299 R + 0.587 G + 0.114 B
    '''
    Grey = np.dot(rgb_image[...,:3],[0.299,0.587,0.114])
    return Grey


#产生DCT变换矩阵
def gen_Matrix(N):
    Matrix = np.zeros((N,N),dtype=float)

    for row in range(N):#遍历行
        if row == 0:#第一行
            for col in range(N):
                Matrix[row,col] = 1 / sqrt(N)
        
        else:
            for col in range(N):
                Matrix[row,col] = sqrt(2/N) * (cos( (np.pi * (2*col+1) * row) / (2*N)))

    return Matrix

#主函数
def main():
    image = mpimg.imread('image3.jpg') #读图像
    Grey = RGB_to_Grey(image) #转为灰度值

    Size = list(Grey.shape)
    Rows,Cols = Size[0],Size[1]

    A = gen_Matrix(Rows)
    B = gen_Matrix(Cols)

    B = B.T

    T = np.dot(np.dot(A,Grey),B)

    T_log = np.log2(T)

    print(T)

    plt.figure(1)
    plt.imshow(Grey,'gray')
    plt.title('Original image')

    plt.figure(2)
    plt.imshow(T_log,'gray')
    plt.title('DCT result')




if __name__ == '__main__':
    main()     

  \space  
仿真结果：(还是那张图)

也不太知道做的对不对，欢迎批评指正。