samples/python/snippets/dft.py#

An example on Discrete Fourier transform (DFT) in python.

  1#!/usr/bin/env python
  2
  3'''
  4sample for discrete fourier transform (dft)
  5
  6USAGE:
  7    dft.py <image_file>
  8'''
  9
 10
 11# Python 2/3 compatibility
 12from __future__ import print_function
 13
 14import numpy as np
 15import cv2 as cv
 16
 17import sys
 18
 19
 20def shift_dft(src, dst=None):
 21    '''
 22        Rearrange the quadrants of Fourier image so that the origin is at
 23        the image center. Swaps quadrant 1 with 3, and 2 with 4.
 24
 25        src and dst arrays must be equal size & type
 26    '''
 27
 28    if dst is None:
 29        dst = np.empty(src.shape, src.dtype)
 30    elif src.shape != dst.shape:
 31        raise ValueError("src and dst must have equal sizes")
 32    elif src.dtype != dst.dtype:
 33        raise TypeError("src and dst must have equal types")
 34
 35    if src is dst:
 36        ret = np.empty(src.shape, src.dtype)
 37    else:
 38        ret = dst
 39
 40    h, w = src.shape[:2]
 41
 42    cx1 = cx2 = w // 2
 43    cy1 = cy2 = h // 2
 44
 45    # if the size is odd, then adjust the bottom/right quadrants
 46    if w % 2 != 0:
 47        cx2 += 1
 48    if h % 2 != 0:
 49        cy2 += 1
 50
 51    # swap quadrants
 52
 53    # swap q1 and q3
 54    ret[h-cy1:, w-cx1:] = src[0:cy1 , 0:cx1 ]   # q1 -> q3
 55    ret[0:cy2 , 0:cx2 ] = src[h-cy2:, w-cx2:]   # q3 -> q1
 56
 57    # swap q2 and q4
 58    ret[0:cy2 , w-cx2:] = src[h-cy2:, 0:cx2 ]   # q2 -> q4
 59    ret[h-cy1:, 0:cx1 ] = src[0:cy1 , w-cx1:]   # q4 -> q2
 60
 61    if src is dst:
 62        dst[:,:] = ret
 63
 64    return dst
 65
 66
 67def main():
 68    if len(sys.argv) > 1:
 69        fname = sys.argv[1]
 70    else:
 71        fname = 'baboon.jpg'
 72        print("usage : python dft.py <image_file>")
 73
 74    im = cv.imread(cv.samples.findFile(fname))
 75
 76    # convert to grayscale
 77    im = cv.cvtColor(im, cv.COLOR_BGR2GRAY)
 78    h, w = im.shape[:2]
 79
 80    realInput = im.astype(np.float64)
 81
 82    # perform an optimally sized dft
 83    dft_M = cv.getOptimalDFTSize(w)
 84    dft_N = cv.getOptimalDFTSize(h)
 85
 86    # copy A to dft_A and pad dft_A with zeros
 87    dft_A = np.zeros((dft_N, dft_M, 2), dtype=np.float64)
 88    dft_A[:h, :w, 0] = realInput
 89
 90    # no need to pad bottom part of dft_A with zeros because of
 91    # use of nonzeroRows parameter in cv.dft()
 92    cv.dft(dft_A, dst=dft_A, nonzeroRows=h)
 93
 94    cv.imshow("win", im)
 95
 96    # Split fourier into real and imaginary parts
 97    image_Re, image_Im = cv.split(dft_A)
 98
 99    # Compute the magnitude of the spectrum Mag = sqrt(Re^2 + Im^2)
100    magnitude = cv.sqrt(image_Re**2.0 + image_Im**2.0)
101
102    # Compute log(1 + Mag)
103    log_spectrum = cv.log(1.0 + magnitude)
104
105    # Rearrange the quadrants of Fourier image so that the origin is at
106    # the image center
107    shift_dft(log_spectrum, log_spectrum)
108
109    # normalize and display the results as rgb
110    cv.normalize(log_spectrum, log_spectrum, 0.0, 1.0, cv.NORM_MINMAX)
111    cv.imshow("magnitude", log_spectrum)
112
113    cv.waitKey(0)
114    print('Done')
115
116
117if __name__ == '__main__':
118    print(__doc__)
119    main()
120    cv.destroyAllWindows()