Python 3D-DCT

from scipy.fftpack import dct,idct
import numpy as np

def dct3(gop):
    '''
    gop shape is bath*h*w
    '''
    assert type(gop) == np.ndarray,"输入必须为numpy对象"
    
    gop_dct = []
    # 先对图像做二维的dct变换
    for frame in gop:
        frame_dct2 = dct(dct(frame.astype(np.double), norm='ortho').T,norm='ortho').T
        gop_dct.append(frame_dct2)
    gop_dct = np.array(gop_dct)
    result = np.zeros_like(gop_dct)
    _,h,w = gop.shape
    
    # 对帧间做dct变换
    for i in range(h):
        for j in range(w):
            result[:,i,j] = dct(gop_dct[:,i,j],norm='ortho')
    return result
 
    
def idct3(gop_dct):
    '''
    gop_dct shape is bath*h*w
    '''
    assert type(gop_dct) == np.ndarray,"输入必须为numpy对象"
    
    # 先对帧间做idct
    _,h,w = gop_dct.shape
    frame_rec = np.zeros_like(gop_dct)
    
    for i in range(h):
        for j in range(w):
            frame_rec[:,i,j] = idct(gop_dct[:,i,j],norm='ortho')
      
    result = []
    for frame in frame_rec:
        result.append(idct(idct(frame, norm='ortho').T,norm='ortho').T)
    return np.array(result)