imutils.py 文件源码

def _shift_fft(array, shift_value):
    Ndim  = array.ndim
    dims  = array.shape
    dtype = array.dtype.kind

    if (dtype != 'f'):
        raise ValueError('Array must be float')

    shifted = array
    if (Ndim == 1):
        Nx = dims[0]

        x_ramp = np.arange(Nx, dtype=array.dtype) - Nx//2

        tilt = (2*np.pi/Nx) * (shift_value[0]*x_ramp)

        cplx_tilt = np.cos(tilt) + 1j*np.sin(tilt)
        cplx_tilt = fft.fftshift(cplx_tilt)
        narray    = fft.fft(fft.ifft(array) * cplx_tilt)
        shifted   = narray.real
    elif (Ndim == 2):
        Nx = dims[0]
        Ny = dims[1]

        x_ramp = np.outer(np.full(Nx, 1.), np.arange(Ny, dtype=array.dtype)) - Nx//2
        y_ramp = np.outer(np.arange(Nx, dtype=array.dtype), np.full(Ny, 1.)) - Ny//2

        tilt = (2*np.pi/Nx) * (shift_value[0]*x_ramp+shift_value[1]*y_ramp)

        cplx_tilt = np.cos(tilt) + 1j*np.sin(tilt)        
        cplx_tilt = fft.fftshift(cplx_tilt)

        narray    = fft.fft2(fft.ifft2(array) * cplx_tilt)
        shifted   = narray.real
    else:
        raise ValueError('This function can shift only 1D or 2D arrays')

    return shifted