import glob
import os 
from time import time, sleep
from math import sqrt 
import cv2 as cv
import numpy as np
from matplotlib import pyplot as plt
 
# load test video
def load_video():
	video = cv.VideoCapture("video_1.mp4")
	return(video)

# load templates
def load_templates(shape):
	folder = shape
	my_path = glob.glob(folder)
	templates = []
	for template in my_path:
		for filename in os.listdir(folder):
			# read template
			template = cv.imread(os.path.join(folder,filename))
			template = cv.resize(template, (50, 50))
			if template is not None:
				filename = filename.replace('.jpg', '')
				filename = filename.replace('_', ' ')
				filename = filename.replace('2', '')
				template = [template, filename]
				templates.append(template)
	return(templates)
	
# load shapes
def load_shapes():
	folder = 'templates/shapes'
	my_path = glob.glob(folder)
	for shape in my_path:
		shapes = []
		for filename in os.listdir(folder):
			# read template
			shape = cv.imread(os.path.join(folder,filename), 0)
			if shape is not None:
				filename = filename.replace('.jpg', '')
				# save template with its name
				shape = [shape, filename]
				shapes.append(shape)
	return(shapes)

# edges detection	
def edge_detection(image):
	
	gray_image = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
	# filter image 
	filtered_image = cv.bilateralFilter(gray_image, 3, 200, 200)
	
	# lower and upper threshold values
	sigma = 0.2
	value = 60
	threshold = auto_canny(sigma, value)
	
	# edge detection							   lower_value	 upper_value
	edge_detected_image = cv.Canny(filtered_image, threshold[0], threshold[1])
	
	# find edge contours
	contours = cv.findContours(edge_detected_image, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)[1]
	
	positions = []
	for contour in contours:
		# ignore too small and large area of contours
		h, w, _ = image.shape
		image_area = w*h
		contour_area = cv.contourArea(contour)
		if (contour_area < image_area * 0.0005) or (contour_area > image_area * 0.02):
			continue
		
		position = cv.boundingRect(contour)	# x, y, w, h
		if (position[2]/position[3] > 1.2) or (position[2]/position[3] < 0.8):
			continue
		
		position = [position, contour]
		positions.append(position)
		
	return(positions)
	
# auto Canny threshold values
def auto_canny(sigma, value):
	
	# lower and upper threshold values
	lower = int(max(0, (1.0 - sigma) * value))
	upper = int(min(255, (1.0 + sigma) * value))
	
	return(lower, upper)
	
# threshold detection
def thresh_detection(image):	
	
	# convert color to gray
	gray_img = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
	
	# thresholding
	#thresh, thresh_image = cv.threshold(gray_img, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)		# get auto thrash value (otsu, triangle)
	thresh, thresh_image = cv.threshold(gray_img, 75, 255, cv.THRESH_BINARY_INV)
	#thresh_image = cv.adaptiveThreshold(gray_img, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY_INV, 15, 1)
	
	# find contours
	_, contours, _ = cv.findContours(thresh_image, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
	
	positions = []
	for contour in contours:
		# ignore too small and large area of contours
		h, w, _ = image.shape
		image_area = w*h
		contour_area = cv.contourArea(contour)
		if (contour_area < image_area * 0.0005) or (contour_area > image_area * 0.02):
			continue
		
		position = cv.boundingRect(contour)	# x, y, w, h
		position = [position, contour]
		positions.append(position)
		
	return(positions)
	
def color_detection(image):

	# convert to HSV
	image_HSV = cv.cvtColor(image, cv.COLOR_BGR2HSV)	
			
	# red
	H_red = 175
	S_red = 158
	V_red = 118
	
	v1_red = 20
	v2_red = 60
	v3_red = 120
	
	#blue
	H_blue = 110
	S_blue = 182
	V_blue = 63
	
	v1_blue = 15
	v2_blue = 60
	v3_blue = 120
	
	#white
	H_white = 40
	S_white = 14
	V_white = 143
	
	v1_white = 65
	v2_white = 60
	v3_white = 30
		
	# Threshold values
	lower_red = np.array([H_red-v1_red, S_red-v2_red, V_red-v3_red])
	upper_red = np.array([H_red+v1_red, S_red+v2_red, V_red+v3_red])
		
	lower_blue = np.array([H_blue-v1_blue, S_blue-v2_blue, V_blue-v3_blue])
	upper_blue = np.array([H_blue+v1_blue, S_blue+v2_blue, V_blue+v3_blue])
	
	lower_white = np.array([H_white-v1_white, S_white-v2_white, V_white-v3_white])
	upper_white = np.array([H_white+v1_white, S_white+v2_white, V_white+v3_white])	
	
	# create mask
	mask_red = cv.inRange(image_HSV, lower_red, upper_red)
	mask_blue = cv.inRange(image_HSV, lower_blue, upper_blue)
	mask_white = cv.inRange(image_HSV, lower_white, upper_white)
	
	mask = np.array([mask_red, mask_blue, mask_white])
		
	for i in range(0,3):
	
		mask_color = mask[i]
		
		# find color contours
		_, contours, _= cv.findContours(mask_color, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
		
		positions = []
		for contour in contours:
			# ignore too small and large area of contours
			h, w, _ = image.shape
			image_area = w*h
			contour_area = cv.contourArea(contour)
			if (contour_area < image_area * 0.0005) or (contour_area > image_area * 0.02):
				continue
			
			position = cv.boundingRect(contour)	# x, y, w, h
			if (position[2]/position[3] > 1.2) or (position[2]/position[3] < 0.8):
				continue
			
			position = [position, contour]
			positions.append(position)
		
	return(positions)
	
def shape_detection(position_values):
	shapes = []
	for position_value in position_values:
		# get detected contours
		contour = position_value[1]
		
		# get position of contours
		position = position_value[0]
		
		# get contours to zero value
		contour = contour - [position[0], position[1]]
		
		# create detected shape
		# create white matrix
		shape = np.zeros((position[2], position[3]), dtype=np.uint8) + 255
		# draw black contours to white matrix
		cv.drawContours(shape, [contour], 0, (0, 0, 0), -1)
		
		# resize detected shape (template size - 50x50)
		shape = cv.resize(shape, (50, 50))
		
		# load shape templates
		shape_templates = load_shapes()
		for shape_template in shape_templates:
			# get shape name
			shape_name = shape_template[1]
			# get shape template
			shape_template = shape_template[0]
			
			# correlation
			shape = shape - np.mean(shape)
			shape_template = shape_template - np.mean(shape_template)
			
			SP = np.sum(shape * shape_template)
			SSx_2 = np.sum(shape * shape)
			SSy_2 = np.sum(shape_template * shape_template)
		
			max_corr = SP/(sqrt(SSx_2) * sqrt(SSy_2))
			
			# if correlation value > 70%....shape recognized 
			if max_corr**2 > 0.8**2:
				# save detected shape
				det_shape_name = shape_name	
				det_shape = [position, det_shape_name]
				shapes.append(det_shape)
		
		'''
		# approximate contours 
		approx = cv.approxPolyDP(contour, 0.017 * cv.arcLength(contour, True), True)
	
		# find shapes of detected contours
		if len(approx) == 3:
			shape = 'triangle'
		elif len(approx) == 4:
			shape = 'square'
		elif len(approx) > 8:
			shape = 'circle'
		else:
			shape = ''
	
		# find positions of detected shape
		if shape:
			position = cv.boundingRect(contour)	# positions = [x, y, w, h]
			det_shape = [position, shape]
			shapes.append(det_shape)
			'''
			
	# return [position, shape]		
	return(shapes)

# main
video = load_video()
codec = cv.VideoWriter_fourcc(*'XVID')
vw = cv.VideoWriter('vytvorene_video.avi', codec, 30, (960, 540))


while True:
	
	# start measure time for recognition
	start = time()
	
	_, image = video.read()	
	image = cv.resize(image, (int(image.shape[1]/2), int(image.shape[0]/2)))
	
	# use detection method and get detected contours and positions 
	#positions = edge_detection(image)
	#positions = thresh_detection(image)
	positions = color_detection(image)
	
	# recognize shapes
	# [position, shape]
	shape_values = shape_detection(positions)
	
	# detect and recognize signs 
	detected = False
	sign_name = ''
	
	for shape_value in shape_values:
	
		# get shape
		shape = shape_value[1]
		
		# get positions of detected shapes
		position = shape_value[0]
		x = position[0]	# x-position
		y = position[1]	# y-position
		w = position[2]	# width
		h = position[3]	# heigh
		
		# get detected area in image
		det_image = image[y:y+h, x:x+w]
		
		# draw detected area into image
		cv.rectangle(image, (x, y), (x+w, y+h), (100, 100, 255), 3)
		
		# convert detected image color to gray
		det_image = cv.cvtColor(det_image, cv.COLOR_BGR2GRAY)					############################
		
		# resize detected image (template size - 50x50)
		det_image = cv.resize(det_image, (50, 50))
		
		# load templates of detected shape
		templates = load_templates(shape)
		
		# set threshold values of correlation 
		if shape in ['triangle_rot', 'rhombus']:
			best_corr = 0.05
			best_corr_2 = 0.1
			best_corr_3 = 0.5
			best_corr_4 = 0.1
		elif shape in ['octagon']:
			best_corr = 0.97
			best_corr_2 = 0.97
			best_corr_3 = 0.98
			best_corr_4 = 0.99
		else:		
			best_corr = 0.4
			best_corr_2 = 0.4
			best_corr_3 = 0.92
			best_corr_4 = 0.9999
		
		detected = False
		for template in templates:
			# read template and its name
			template_name = template[1]
			template = template[0]
			
			# convert template color to gray							######################################
			template = cv.cvtColor(template, cv.COLOR_BGR2GRAY)
			'''
			# My correlation method
			det_image = det_image - np.mean(det_image)
			template = template - np.mean(template)
			
			SP = np.sum(det_image * template)
			SSx_2 = np.sum(det_image * det_image)
			SSy_2 = np.sum(template * template)
			
			# corelation value	
			max_corr = SP/(sqrt(SSx_2) * sqrt(SSy_2))
			
			# compare threshold value and correlation value
			if max_corr**2 > best_corr**2:
				# save new threshold value
				best_corr = max_corr
				if template_name in sign_name:
					continue
				# save recognized sign name
				sign_name = template_name
				method = 'My method:'
				detected = True
			
			# Method: cv.TM_CCOEFF_NORMED
			res = cv.matchTemplate(det_image, template, cv.TM_CCOEFF_NORMED)
			min_corr, max_corr, min_loc, max_loc = cv.minMaxLoc(res)
			
			# compare threshold value and correlation value
			if max_corr**2 > best_corr_2**2:
				# save new threshold value
				best_corr_2 = max_corr
				if template_name in sign_name:
					continue
				# save recognized sign name
				sign_name = template_name
				method = 'cv.TM_CCOEFF_NORMED:'
				detected = True
			'''
			# Method: cv.TM_CCORR_NORMED
			res = cv.matchTemplate(det_image, template, cv.TM_CCORR_NORMED)
			min_corr, max_corr, min_loc, max_loc = cv.minMaxLoc(res)
			
			# compare threshold value and correlation value
			if max_corr**2 > best_corr_3**2:
				# save new threshold value
				best_corr_3 = max_corr
				if template_name in sign_name:
					continue
				# save recognized sign name
				sign_name = template_name
				method = 'cv.TM_CCORR_NORMED:'
				detected = True
			'''
			# Method: cv.TM_SQDIFF_NORMED
			res = cv.matchTemplate(det_image, template, cv.TM_SQDIFF_NORMED)
			min_corr, max_corr, min_loc, max_loc = cv.minMaxLoc(res)
			
			# compare threshold value and correlation value
			if max_corr**2 > best_corr_4**2:
				# save new threshold value
				best_corr_4 = max_corr
				if template_name in sign_name:
					continue
				# save recognized sign name
				sign_name = template_name
				method = 'cv.TM_CCOEFF_NORMED:'
				detected = True
			'''
		
		if detected:
			
			print('Shoda: ', sign_name, best_corr_3*100, '%')
			cv.putText(image, sign_name, (x, y-10), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2, cv.LINE_4) 
			cv.rectangle(image, (x, y), (x+w, y+h), (0, 255, 255), 6)
	
	# finish measure time
	end = time()
	#print('cas detekce: ', (end - start) * 1000, '[ms]')
				
	# show results
	cv.imshow("video", image)
	# save video
	vw.write(image)
		
	if cv.waitKey(1) & 0xFF == ord("x"):
		break

vw.release()
video.release()
cv.destroyAllWindows()

