# ================================================================
#   Copyright (C) 2021 GXHL Ltd. All rights reserved.
#
#   Author      : yke
#   Created date: 2021-03-19
#
# ================================================================

import cv2
import time
import numpy as np
from tensorflow._api.v2.compat.v1 import Session, ConfigProto, keras
import core.utils as utils
import tensorflow as tf
from core.yolov3 import YOLOv3, decode
from getopt import getopt
import sys

#################### 命令参数输入 start ####################
inputFile = ""
outputFile = ""
writeVideo = False
## 解析命令行参数
try:
   opts, args = getopt(sys.argv[1:], "hi:o:", ["inputFile=","outputFile="])
except getopt.GetoptError:
   print("video.py -i <inputFile> -o <outputFile>")
   sys.exit(2)
else:
   for opt, arg in opts:
       if opt == "-h":
          print("video.py -i <inputFile> -o <outputFile>")
          sys.exit()
       elif opt in ("-i", "--inputFile"):
            inputFile = arg
       elif opt in ("-o", "--outputFile"):
            outputFile = arg
            writeVideo = True
            if not outputFile.endswith(".mp4"):
               outputFile += ".mp4"
## 没有指定命令行参数的时候设置默认值
if outputFile == "" and inputFile == "" :
   outputFile = ".\output.mp4"
   writeVideo = True
if inputFile == "":
   inputFile = input("Input video filename:")
## 后期输入依然没有值的时候使用本机默认摄像头抓拍
if inputFile == "":
   video_path = 0
else:
   video_path = inputFile
## 打开文件检查内容
vid = cv2.VideoCapture(video_path)
return_value, frame = vid.read()
if return_value:
   frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
else:
   print("video no content!")
   sys.exit()
#################### 命令参数输入 ended ####################

#################### 应用参数配置 start ####################
config = ConfigProto()
## 应对显存不够
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = 0.5
keras.backend.set_session(Session(config = config))

num_classes = 80
## 加载model
bbox_tensors = []
input_size = 416
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
for i, fm in enumerate(YOLOv3(input_layer)):
    bbox_tensor = decode(fm, i)
    bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights(model, "data/yolov3.weights")
model.summary()
#################### 应用参数配置 ended ####################

#################### 鼠标划线配置 start ####################
frame_no = 0
drawing_line = False
point1 = (1, 1)
point = (1, 1)

def draw_line(event, x, y, flags, param):
    '''捕捉鼠标事件用来画线，以用来计数'''
    global drawing_line, point, point1
    if drawing_line:
       point = (x, y)
    if event == cv2.EVENT_LBUTTONDOWN: # 鼠标左键用来划线，不画框
       if drawing_line: # 鼠标左键按下第二次完成划线
          utils.line = [point1, point]
          utils.horizontal_True_vertical_False = utils.assess_horizontal_or_vertical(utils.line)
          drawing_line = False
       else:
          drawing_line = True
          point1 = (x, y)

cv2.namedWindow('result', cv2.WINDOW_AUTOSIZE)
cv2.setMouseCallback('result', draw_line)
#################### 鼠标划线配置 ended ####################

#################### 应用本体代码 start ####################
if writeVideo:
   w, h = int(vid.get(3)), int(vid.get(4))
   fourcc = cv2.VideoWriter_fourcc(*'mp4v')
   video_writer = cv2.VideoWriter(outputFile, fourcc, 24, (w, h))

while return_value:
   if drawing_line:
      # 画线过程中既要保持暂停，又要确保动画效果
      image = np.copy(frame)
      cv2.circle(image, point, 5, (0, 255, 0), 2)
      cv2.circle(image, point1, 5, (0, 255, 0), 2)
      cv2.line(image, point1, point, (255, 0, 0), 2)
      # 这里开新窗口主要是防止有些人手欠，在画线的过程中关窗口
      # cv2.namedWindow("result", cv2.WINDOW_AUTOSIZE)
      # cv2.setMouseCallback('result', draw_line)
      # 因此需要确保重开窗口之后仍然能够继续画线
      result = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
      cv2.imshow("result", result)
      cv2.waitKey(delay=50)
   else:
      # 不画线的过程中正常解析，第一帧的内容已在打开文件时读取进来
      prev_time = time.time()
      frame_size = frame.shape[:2]
      image_data = utils.image_preporcess(np.copy(frame), [input_size, input_size])
      image_data = image_data[np.newaxis, ...].astype(np.float32)
      frame_no += 1
      pred_bbox = model.predict_on_batch(image_data)
      pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]
      pred_bbox = tf.concat(pred_bbox, axis=0)
      bboxes = utils.postprocess_boxes(pred_bbox, frame_size, input_size, 0.25)
      bboxes = utils.nms(bboxes, 0.25, method = 'nms')
      # image = utils.draw_bbox(frame, bboxes)
      #########################################
      curr_time = time.time()
      exec_time = curr_time - prev_time
      fps = 1 / exec_time
      image = utils.video_draw_bbox(frame, bboxes, fps)
      # result = np.asarray(image)
      # cv2.putText(result, text=fps, org=(50, 70), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(255, 0, 0), thickness=2)
      ########################################
      # FPS LOG记录
      # print("yolo_timecost {} {}\n".format(frame_no, exec_time*1000)) # python video_demo.py | grep yolo_timecost >  yolo.log
      # cv2.namedWindow("result", cv2.WINDOW_AUTOSIZE)
      # cv2.setMouseCallback('result', draw_line)
      result = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
      cv2.imshow("result", result)
      #### 录制视频 ####
      ## 注意 目录output必须存在
      if writeVideo:
         video_writer.write(result)
      if cv2.waitKey(1) & 0xFF == ord('q'):
         break # 手动按 q 退出
      ## 处理完上一帧，加载下一帧
      return_value, frame = vid.read()
      if return_value:
         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
      else:
         print("No image!" " video ended.")
         break

if writeVideo:
   video_writer.release()
cv2.destroyAllWindows()
#################### 应用本体代码 ended ####################