Python numpy.fft 模块，fft2() 实例源码

def cross_corr(img1,img2,mask=None):
    '''Compute the autocorrelation of two images.
        Right now does not take mask into account.
        todo: take mask into account (requires extra calculations)
        input: 
            img1: first image
            img2: second image
            mask: a mask array
        output:
            the autocorrelation of the two images (same shape as the correlated images)

    '''
    #if(mask is not None):
    #   img1 *= mask
    #  img2 *= mask

    #img1_mean = np.mean( img1.flat )
    #img2_mean = np.mean( img2.flat )

    # imgc = fftshift( ifft2(     
    #        fft2(img1/img1_mean -1.0 )*np.conj(fft2( img2/img2_mean -1.0 ))).real )

    #imgc = fftshift( ifft2(     
    #        fft2(  img1/img1_mean  )*np.conj(fft2(  img2/img2_mean   ))).real )

    imgc = fftshift( ifft2(     
            fft2(  img1  )*np.conj(fft2(  img2  ))).real )

    #imgc /= (img1.shape[0]*img1.shape[1])**2
    if(mask is not None):
        maskc = cross_corr(mask,mask)        
        imgc /= np.maximum( 1, maskc )


    return imgc

def extractFeatures(self, img):
        "A vector of 2n elements where n is the number of theta angles"
        "and 2 is the number of frequencies under consideration"
        filters =  self.build_filters(img.shape[0],img.shape[1],5,(0.75,1.5),2,1,np.pi/2.0)
        fft_filters = [np.fft.fft2(i) for i in filters]
        img_fft = np.fft.fft2(img)
        a =  img_fft * fft_filters
        s = [np.fft.ifft2(i) for i in a]
        k = [p.real for p in s]
        return k

def fft_convolve2d(self, x,y):
        """ 2D convolution, using FFT"""
        fr = fft2(x)
        fr2 = fft2(np.flipud(np.fliplr(y)))
        m,n = fr.shape
        cc = np.real(ifft2(fr*fr2))
        cc = np.roll(cc, -m/2+1,axis=0)
        cc = np.roll(cc, -n/2+1,axis=1)
        return cc

def xcorr(imageA, imageB):
    FimageA = _fft.fft2(imageA)
    CFimageB = _np.conj(_fft.fft2(imageB))
    return _fft.fftshift(_np.real(_fft.ifft2((FimageA * CFimageB)))) / _np.sqrt(imageA.size)

def blkWlDire(img):
    """Calculate wavelength and direction given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    dire=np.arctan2(origin[0]-mmax[0][0],origin[1]-mmax[1][0])
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl,dire

def blkwl(img):
    """Calculate wavelength  given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl

def blkWlDire(img):
    """Calculate wavelength and direction given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    dire=np.arctan2(origin[0]-mmax[0][0],origin[1]-mmax[1][0])
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl,dire

def blkwl(img):
    """Calculate wavelength  given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl

def blkWlDire(img):
    """Calculate wavelength and direction given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    dire=np.arctan2(origin[0]-mmax[0][0],origin[1]-mmax[1][0])
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl,dire

def blkWlDire(img):
    """Calculate wavelength and direction given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    dire=np.arctan2(origin[0]-mmax[0][0],origin[1]-mmax[1][0])
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl,dire

def blkwl(img):
    """Calculate wavelength  given an image block"""
    f=np.abs(fftshift(fft2(img)))
    origin=np.where(f==np.max(f));f[origin]=0;mmax=np.where(f==np.max(f))
    wl=2*img.shape[0]/(((origin[0]-mmax[0][0])*2)**2+((origin[1]-mmax[1][0])*2)**2)**0.5
    return wl

def test_expand(): 
    X = torch.randn(2,2,4,4).cuda().double()
    zeros = torch.zeros(2,2,4,4).cuda().double()
    r1, r2 = cfft.rfft2(X)
    c1, c2 = cfft.fft2(X, zeros)
    assert np.allclose(cfft.expand(r1).cpu().numpy(), c1.cpu().numpy())
    assert np.allclose(cfft.expand(r2, imag=True).cpu().numpy(), c2.cpu().numpy())
    r1, r2 = cfft.rfft3(X)
    c1, c2 = cfft.fft3(X, zeros)
    assert np.allclose(cfft.expand(r1).cpu().numpy(), c1.cpu().numpy())
    assert np.allclose(cfft.expand(r2, imag=True).cpu().numpy(), c2.cpu().numpy())

    X = torch.randn(2,2,5,5).cuda().double()
    zeros = torch.zeros(2,2,5,5).cuda().double()
    r1, r2 = cfft.rfft3(X)
    c1, c2 = cfft.fft3(X, zeros)
    assert np.allclose(cfft.expand(r1, odd=True).cpu().numpy(), c1.cpu().numpy())
    assert np.allclose(cfft.expand(r2, imag=True, odd=True).cpu().numpy(), c2.cpu().numpy())

def ft2(g, delta):
    return fftshift(fft2(fftshift(g))) * (delta ** 2)