import numpy as np
import cv2
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import pickle
%matplotlib inline
import os
os.chdir('D:\Github\Advanced_Lane_Finding')
# Read in an image
image = mpimg.imread('signs_vehicles_xygrad.jpg')


def abs_sobel_thresh(img, orient='x', sobel_kernel=3, thresh=(0, 255)):
    # Calculate directional gradient
    # Apply threshold
    gray = cv2.cvtColor(img,cv2.COLOR_RGB2GRAY)
    if orient == 'x':
        x = 1
        y = 0
    elif orient == 'y':
        x = 0
        y = 1
    sobelxy = cv2.Sobel(gray,cv2.CV_64F,x,y)
    abs_sobelxy = np.absolute(sobelxy)
    scaled_sobel = np.int8(255*abs_sobelxy/np.max(abs_sobelxy))
    sxbinary = np.zeros_like(scaled_sobel)
    thresh_min = thresh[0]
    thresh_max = thresh[1]
    sxbinary[(scaled_sobel >= thresh_min) & (scaled_sobel <= thresh_max)] = 1
    return sxbinary

def mag_thresh(image, sobel_kernel=3, mag_thresh=(0, 255)):
    # Calculate gradient magnitude
    # Apply threshold

    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize = sobel_kernel)
    sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize = sobel_kernel)
    mag_sobelxy = np.sqrt(np.power(sobelx,2)+np.power(sobely,2))
    scaled_sobelxy = np.uint8(255*mag_sobelxy/np.max(mag_sobelxy))
    thres_min = mag_thresh[0]
    thres_max = mag_thresh[1]
    sbinary = np.zeros_like(scaled_sobelxy)
    sbinary[(scaled_sobelxy >= thres_min) & (scaled_sobelxy <= thres_max)] = 1
    return sbinary

def dir_threshold(image, sobel_kernel=3, thresh=(0, np.pi/2)):
    # Calculate gradient direction
    # Apply threshold
    # Apply the following steps to img
    # 1) Convert to grayscale
    # 2) Take the gradient in x and y separately
    # 3) Take the absolute value of the x and y gradients
    # 4) Use np.arctan2(abs_sobely, abs_sobelx) to calculate the direction of the gradient
    # 5) Create a binary mask where direction thresholds are met
    # 6) Return this mask as your binary_output image
    gray = cv2.cvtColor(image,cv2.COLOR_RGB2GRAY)
    sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=sobel_kernel)
    sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=sobel_kernel)
    abs_sobelx = np.absolute(sobelx)
    abs_sobely = np.absolute(sobely)
    dir_gradient = np.arctan2(abs_sobely, abs_sobelx)
    binary = np.zeros_like(dir_gradient)
    min = thresh[0]
    max = thresh[1]
    binary[(dir_gradient >=  min) & (dir_gradient <= max)] = 1

    return binary


def b_threshold(im, thresh=(30,255)):
    Lab = cv2.cvtColor(im,cv2.COLOR_RGB2Lab)
    b_channel = Lab[:,:,2]
    b_binary = np.zeros_like(b_channel)
    b_binary[(b_channel >= thresh[0]) & (b_channel <= thresh[1])] = 1
    #  b_binary = np.dstack((np.zeros_like(b_binary),np.zeros_like(b_binary),b_binary))
    return b_binary



#%% b_channel
import glob
print(glob.glob('*'))
b_image = mpimg.imread('signs_vehicles_xygrad.jpg')
plt.imshow(cv2.cvtColor(b_image,cv2.COLOR_BGR2Lab),cmap='gray')
plt.imshow(b_threshold(b_image),cmap='gray') # There is no yellow on image


# %% Choose a Sobel kernel size
ksize = 3 # Choose a larger odd number to smooth gradient measurements

# Apply each of the thresholding functions
gradx = abs_sobel_thresh(image, orient='x', sobel_kernel=ksize, thresh=(10, 120))
grady = abs_sobel_thresh(image, orient='y', sobel_kernel=ksize, thresh=(10, 100))
mag_binary = mag_thresh(image, sobel_kernel=ksize, mag_thresh=(20, 100))
dir_binary = dir_threshold(image, sobel_kernel=ksize, thresh=(0.7, 1.3))

combined = np.zeros_like(dir_binary)
combined[((gradx == 1) & (grady == 1)) | ((mag_binary == 1) & (dir_binary == 1))] = 1

f, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))
f.tight_layout()
ax1.imshow(image)
ax1.set_title('Original Image', fontsize=15)
ax2.imshow(combined, cmap='gray')
ax2.set_title('Thresholded Gradient', fontsize=15)
plt.subplots_adjust(left=0., right=1, top=0.9, bottom=0.)

# %% sobel

# Run the function
grad_binary = abs_sobel_thresh(image, orient='x', thresh=(20,100))
# Plot the result
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))
f.tight_layout()
ax1.imshow(image)
ax1.set_title('Original Image', fontsize=15)
ax2.imshow(grad_binary, cmap='gray')
ax2.set_title('Thresholded Gradient', fontsize=15)
plt.subplots_adjust(left=0., right=1, top=0.9, bottom=0.)

# %% mag
# Run the function
mag_binary = mag_thresh(image, sobel_kernel=3, mag_thresh=(30, 100))
# Plot the result
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))
f.tight_layout()
ax1.imshow(image)
ax1.set_title('Original Image', fontsize=15)
ax2.imshow(mag_binary, cmap='gray')
ax2.set_title('Thresholded Magnitude', fontsize=15)
plt.subplots_adjust(left=0., right=1, top=0.9, bottom=0.)

# %% direction
# Run the function
dir_binary = dir_threshold(image, sobel_kernel=15, thresh=(0.7, 1.3))
# Plot the result
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))
f.tight_layout()
ax1.imshow(image)
ax1.set_title('Original Image', fontsize=15)
ax2.imshow(dir_binary, cmap='gray')
ax2.set_title('Thresholded Grad. Dir.', fontsize=15)
plt.subplots_adjust(left=0., right=1, top=0.9, bottom=0.)
