#python imgai.py --image [inputimagepath] --output [outputimagepath]

from imageai.Detection import ObjectDetection
import os
import cv2
import numpy as np
import argparse
import time

ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", required=True,
	help="path to input image")
ap.add_argument("-o", "--output", required=True,
	help="path to output image")
args = vars(ap.parse_args())

execution_path = os.getcwd()

detector = ObjectDetection()
detector.setModelTypeAsRetinaNet()
detector.setModelPath( os.path.join(execution_path , "resnet50_coco_best_v2.0.1.h5"))
detector.loadModel()

#image = cv2.imread("queue-11.jpg")
#(H, W) = image.shape[:2]
#crop_img1 = image[170:300, 0:W-120]
#cv2.imwrite("queue-111.jpg",crop_img1)

custom_objects = detector.CustomObjects(person=True)
#detections = detector.detectCustomObjectsFromImage(custom_objects=custom_objects, input_image=os.path.join(args["image"]), output_image_path=os.path.join(execution_path, "image3custom.jpg"), minimum_percentage_probability=30, display_percentage_probability=False, display_object_name=False)
detections = detector.detectCustomObjectsFromImage(custom_objects=custom_objects, input_image=os.path.join(args["image"]), output_image_path=os.path.join(args["output"]), minimum_percentage_probability=30, display_percentage_probability=True, display_object_name=False)

image = cv2.imread(args["image"])
print(image.shape)
#image = cv2.resize(image, (450,550))
(H, W) = image.shape[:2]



# define lane height
#Lane1_y_start = 425
#Lane1_y_end = 600
#Lane2_y_start = 300
#Lane2_y_end = 480
#Lane3_y_start = 170
#Lane3_y_end = 300
#Lane1_x_start = 0
#Lane1_x_end = W
#Lane2_x_start = 0
#Lane2_x_end = W
#Lane3_x_start = 0
#Lane3_x_end = W-120

lineX = [[672,164,213,796],
		 [867,125,483,834],
		 [1151,168,1527,876],
		 [1386,197,1826,808]]
#lineX = [[170,H - 230,590,H - 890], [497,H - 194,744,H - 890]]



#lane=[[0,0,0,0,0,0,0,0], [0,0,0,0,0,0,0,0]]


aa=[]
counter=[0,0,0,0]
#counter=[0,0]
#crop_img1 = image[425:600, 0:W]
#crop_img2 = image[300:480, 0:W]
#crop_img3 = image[170:300, 0:W-120]
image = cv2.imread(args["output"])
for eachObject in detections:
	a = 0

	lineX2 = [[0, 0],
			  [0, 0],
			  [0, 0],
			  [0, 0]]
	#print(eachObject["name"] , " : ", eachObject["percentage_probability"], " : ", eachObject["box_points"] )
	#print("--------------------------------")
	# check if the rectangle is within lane 1 or lane 2 or lane 3
	#if ((eachObject["box_points"][1] >= Lane1_y_start) and (eachObject["box_points"][3] <= Lane1_y_end) and (eachObject["box_points"][0] >= Lane1_x_start) and (eachObject["box_points"][2] <= Lane1_x_end)):
		#counter1 += 1
	#elif ((eachObject["box_points"][1] > Lane2_y_start) and (eachObject["box_points"][3] < Lane2_y_end) and (eachObject["box_points"][0] >= Lane2_x_start) and (eachObject["box_points"][2] <= Lane2_x_end)):
		#counter2 += 1
	#elif ((eachObject["box_points"][1] > Lane3_y_start) and (eachObject["box_points"][3] < Lane3_y_end) and (eachObject["box_points"][0] >= Lane3_x_start) and (eachObject["box_points"][2] <= Lane3_x_end)):
		#counter3 += 1
	aa=[int((eachObject["box_points"][2]-eachObject["box_points"][0])/2)+eachObject["box_points"][0],eachObject["box_points"][1],int((eachObject["box_points"][2]-eachObject["box_points"][0])/2)+eachObject["box_points"][0],eachObject["box_points"][3]]

	if (aa[3]-aa[1])<150:
		aa[3]=aa[1]+180


	#print(aa)
	#print(int((eachObject["box_points"][2]-eachObject["box_points"][0])/2))

	for j in lineX2:

		j[0] = ((aa[0] - aa[0]) * (
					lineX[a][1] - aa[1]) - (
						aa[3] - aa[1]) * (
							lineX[a][0] - aa[0])) / (
					   (aa[3] - aa[1]) * (lineX[a][2] - lineX[a][0]) - (
					   aa[0] - aa[0]) * (lineX[a][3] - lineX[a][1]))

		j[1] = ((lineX[a][2] - lineX[a][0]) * (lineX[a][1] - aa[1]) - (
					lineX[a][3] - lineX[a][1]) * (
						lineX[a][0] - aa[0])) / (
					   (aa[3] - aa[1]) * (lineX[a][2] - lineX[a][0]) - (
					   aa[0] - aa[0]) * (lineX[a][3] - lineX[a][1]))



		if (0 <= j[0] <= 1 and 0 <= j[1] <= 1):
			counter[a] += 1
			#cv2.line(image, (aa[0], aa[1]), (aa[2], aa[3]), (255, 255, 0), 2)
			#print(eachObject["percentage_probability"], " : ", eachObject["box_points"])
		a+=1

#image = cv2.imread(args["output"])

text2='Lane 1 : %d, Lane 2 : %d, Lane 3 : %d, Lane 4 : %d' % (counter[0], counter[1], counter[2], counter[3])
cv2.rectangle(image, (10, 10),(650,35),(255, 255, 0),-1)
cv2.putText(image, text2, (10, 30), cv2.FONT_HERSHEY_SIMPLEX,	0.8, (255,0,255), 2)

cv2.line(image, (lineX[0][0], lineX[0][1]), (lineX[0][2], lineX[0][3]), (0,0,255), 2)  # lane 1
cv2.line(image, (lineX[1][0], lineX[1][1]), (lineX[1][2], lineX[1][3]), (0,0,255), 2)  # lane 2
cv2.line(image, (lineX[2][0], lineX[2][1]), (lineX[2][2], lineX[2][3]), (0,0,255), 2)  # lane 3
cv2.line(image, (lineX[3][0], lineX[3][1]), (lineX[3][2], lineX[3][3]), (0,0,255), 2)  # lane 4
cv2.line(image, (int((556-512)/2), 125), (int((556-512)/2), 188), (255, 0, 0), 2)

	#cv2.rectangle(image, (Lane3_x_start, Lane3_y_start), (Lane3_x_end, Lane3_y_end), (0,0,255), 2)  # lane 3
cv2.imwrite(args["output"],image)
cv2.imshow(args["output"],image)
cv2.waitKey(0)