#-*- coding: UTF-8 -*-

import numpy as np
import cv2
from random import randint

# 单目标追踪

# 创建一个VideoCapture对象
cap = cv2.VideoCapture(1)

# 指定视频编解码方式为MJPG
codec = cv2.VideoWriter_fourcc(*'MJPG')
fps = 20.0 # 指定写入帧率为20
frameSize = (380, 380) # 指定窗口大小
# 创建 VideoWriter对象
# out = cv2.VideoWriter('video_record.avi', codec, fps, frameSize)

# Specify the tracker type
trackerType = "CSRT"   

# Create MultiTracker object
multiTracker = cv2.MultiTracker_create()

trackerTypes = ['BOOSTING', 'MIL', 'KCF','TLD', 'MEDIANFLOW', 'GOTURN', 'MOSSE', 'CSRT']

def createTrackerByName(trackerType):
    # Create a tracker based on tracker name
    if trackerType == trackerTypes[0]:
        tracker = cv2.TrackerBoosting_create()
    elif trackerType == trackerTypes[1]: 
        tracker = cv2.TrackerMIL_create()
    elif trackerType == trackerTypes[2]:
        tracker = cv2.TrackerKCF_create()
    elif trackerType == trackerTypes[3]:
        tracker = cv2.TrackerTLD_create()
    elif trackerType == trackerTypes[4]:
        tracker = cv2.TrackerMedianFlow_create()
    elif trackerType == trackerTypes[5]:
        tracker = cv2.TrackerGOTURN_create()
    elif trackerType == trackerTypes[6]:
        tracker = cv2.TrackerMOSSE_create()
    elif trackerType == trackerTypes[7]:
        tracker = cv2.TrackerCSRT_create()
    else:
        tracker = None
        print('Incorrect tracker name')
        print('Available trackers are:')
    for t in trackerTypes:
        print(t)
    return tracker

colors = []
colors.append((randint(0, 255), randint(0, 255), randint(0, 255)))

print("按键Q-结束视频录制")
while(cap.isOpened()):

    ret, frame = cap.read()
    # 如果帧获取正常
    if ret==True:
        # 镜像反转 flip
        # frame = cv2.flip(frame, -1)

        # 不断的向视频输出流写入帧图像
        # out.write(frame)

        # 画框
        if cv2.waitKey(1) == ord('c'):
            bbox = cv2.selectROI('MultiTracker', frame)
            multiTracker.add(createTrackerByName(trackerType), frame, bbox)
            cv2.destroyWindow('MultiTracker')

        # 转换为gray灰度图
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # 寻找轮廓
        bimg, contours, hier = cv2.findContours(gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for cidx,cnt in enumerate(contours):
            (x, y, w, h) = cv2.boundingRect(cnt)
            print('RECT: x={}, y={}, w={}, h={}'.format(x, y, w, h))
            # 原图绘制圆形
            cv2.rectangle(gray, pt1=(x, y), pt2=(x+w, y+h),color=(255, 255, 255), thickness=3)
            # 截取ROI图像
            # cv2.imwrite("./car/my_car_{}.png".format(cidx), gray[y:y+h, x:x+w])


        # get updated location of objects in subsequent frames
        success, boxes = multiTracker.update(frame)

        # draw tracked objects
        for i, newbox in enumerate(boxes):
            p1 = (int(newbox[0]), int(newbox[1]))
            p2 = (int(newbox[0] + newbox[2]), int(newbox[1] + newbox[3]))
            cv2.rectangle(frame, p1, p2, colors[i], 2, 1)


        # 在窗口中展示画面
        cv2.imshow('frame',gray)
        if cv2.waitKey(1) == ord('q'):
            break
    else:
        break

# 释放资源
cap.release()
# out.release()
cv2.destroyAllWindows()