"""
Fourier Transforms

The frequency components of an image can be displayed after doing a Fourier Transform (FT).
An FT looks at the components of an image (edges that are high-frequency, and areas of smooth
color as low-frequency), and plots the frequencies that occur as points in spectrum.

In fact, an FT treats patterns of intensity in an image as sine waves with a particular frequency,
and you can look at an interesting visualization of these sine wave components on this page.

We'll first look at a few simple image patterns to build up an idea of what image frequency
components look like, and then transform a more complex image to see what it looks like in the frequency domain.

"""

import numpy as np
import matplotlib.pyplot as plt
import cv2

# Read in the images
image_stripes = cv2.imread('images/stripes.jpg')
image_solid = cv2.imread('images/pink_solid.jpg')

# Change color to RGB (from BGR)
image_stripes = cv2.cvtColor(image_stripes, cv2.COLOR_BGR2RGB)
image_solid = cv2.cvtColor(image_solid, cv2.COLOR_BGR2RGB)

# Display the images
f, (ax1,ax2) = plt.subplots(1, 2, figsize=(10,5))

ax1.imshow(image_stripes)
ax2.imshow(image_solid)

# convert to grayscale to focus on the intensity patterns in the image
gray_stripes = cv2.cvtColor(image_stripes, cv2.COLOR_RGB2GRAY)
gray_solid = cv2.cvtColor(image_solid, cv2.COLOR_RGB2GRAY)

# normalize the image color values from a range of [0,255] to [0,1] for further processing
norm_stripes = gray_stripes / 255.0
norm_solid = gray_solid / 255.0

# perform a fast fourier transform and create a scaled, frequency transform image
def ft_image(norm_image):
    '''
        This function takes in a normalized, grayscale image
        and returns a frequency spectrum transform of that image.
    '''
    f = np.fft.fft2(norm_image)
    fshift = np.fft.fftshift(f)
    frequency_tx = 20 * np.log(np.abs(fshift))

    return frequency_tx

f_stripes = ft_image(norm_stripes)
f_solid = ft_image(norm_solid)

# display the images
# original images to the left of their frequency transform
f, (ax1,ax2,ax3,ax4) = plt.subplots(1, 4, figsize=(20,10))

ax1.set_title('original image')
ax1.imshow(image_stripes)
ax2.set_title('frequency transform image')
ax2.imshow(f_stripes, cmap='gray')

ax3.set_title('original image')
ax3.imshow(image_solid)
ax4.set_title('frequency transform image')
ax4.imshow(f_solid, cmap='gray')

"""
Low frequencies are at the center of the frequency transform image.
The transform images for these example show that the solid image has most 
low-frequency components (as seen by the center bright spot).

The stripes transform image contains low-frequencies for the areas of white 
and black color and high frequencies for the edges in between those colors. 
The stripes transform image also tells us that there is one dominating direction 
for these frequencies; vertical stripes are represented by a horizontal line passing 
through the center of the frequency transform image.
"""

# Read in an image
image = cv2.imread('images/birds.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
norm_image = gray/255.0

f_image = ft_image(norm_image)
f, (ax1,ax2) = plt.subplots(1, 2, figsize=(20,10))

ax1.imshow(image)
ax2.imshow(f_image, cmap='gray')
"""
This image has components of all frequencies. You can see a bright spot in the center 
of the transform image, which tells us that a large portion of the image is low-frequency; 
this makes sense since the body of the birds and background are solid colors. The transform 
image also tells us that there are two dominating directions for these frequencies; 
vertical edges (from the edges of birds) are represented by a horizontal line passing through 
the center of the frequency transform image, and horizontal edges (from the branch and tops 
of the birds' heads) are represented by a vertical line passing through the center.
"""