import cv2
import os
import numpy as np
from ultralytics import YOLO
from collections import defaultdict

model=YOLO('./Model/yolov8n.pt')
VIDEO_PATH='./Test_Video/person_video.mp4'
RESULT_PATH='result0.mp4'

#记录所有id的位置信息
track_history=defaultdict(lambda :[])

#指定多边形的顶点坐标
polygonPoints=np.array([[410,150],[810,150],[510,300],[110,300]])

def trace_check(VIDEO_PATH):
    if not os.path.exists(VIDEO_PATH):
        print(f'视频文件不存在: {VIDEO_PATH}')
        exit()

    capture = cv2.VideoCapture(VIDEO_PATH)
    if not capture.isOpened():
        print('无法打开视频文件')
        exit()


    fps=int(capture.get(cv2.CAP_PROP_FPS))                          #获取帧率
    frame_width=int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))     #获取宽度
    frame_height=int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))   #获取高度
    print(f'帧率:{fps},宽度:{frame_width},高度:{frame_height}')

    videoWriter=None

    while True:
        success,frame=capture.read()     #读取视频中的一帧
        if not success:
            print('视频读取完成')
            break

        results=model.track(frame,persist=True)
        #可视化显示
        a_frame=results[0].plot()

        #装的是所有id的位置信息
        boxes=results[0].boxes.xywh.cpu()

        #装的是所有id的序列号信息
        track_ids=results[0].boxes.id.int().cpu().tolist()
        for box,track_id in zip(boxes,track_ids):
            x,y,w,h=box
            track=track_history[track_id]
            track.append((float(x),float(y)))

            #最多为50个点，超过会把第一个点清除
            if len(track)>50:
                track.pop(0)

            #当前的track_id所有经过的轨迹路径点(不超过50个)
            points=np.hstack(track).astype(np.int32).reshape(-1,1,2)
            cv2.polylines(a_frame,[points],isClosed=False,color=(255,0,255),thickness=1)

            # 保存处理好的视频
            if videoWriter is None:
                fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
                videoWriter = cv2.VideoWriter(RESULT_PATH, fourcc, fps, (int(frame_width), int(frame_height)))
            videoWriter.write(a_frame)

        cv2.imshow(winname='window',mat=a_frame)

        #按空格键退出
        if cv2.waitKey(1) & 0xFF == ord(' '):
            break

    #释放内存资源
    capture.release()
    videoWriter.release()
    cv2.destroyAllWindows()

trace_check(VIDEO_PATH)