import cv2
from ultralytics import YOLO,solutions
from shapely.geometry import LineString, Point, Polygon
import numpy as np

model=YOLO('./Model/yolov8n.pt')
VIDEO_PATH='./Test_Video/person_video.mp4'
RESULT_PATH='result4.mp4'
capture=cv2.VideoCapture(VIDEO_PATH)

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

w,h,fps=(int(capture.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH,cv2.CAP_PROP_FRAME_HEIGHT,cv2.CAP_PROP_FPS))

#画一个多边形
polygonPoints=[(410,150),(810,150),(510,300),(110,300)]

#处理视频并且定义编码格式
video_write=cv2.VideoWriter(RESULT_PATH,cv2.VideoWriter_fourcc(*'mp4v'),fps,(w,h))

#声明对象计数函数
counter=solutions.ObjectCounter(
    view_img=True,
    reg_pts=polygonPoints,
    names=model.names,
    draw_tracks=True,
    line_thickness=1,
)

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

    #在每一帧中绘制多边形
    cv2.polylines(frame,pts=[np.array(polygonPoints)],isClosed=True,color=(0,0,255),thickness=2)
    # 生成一张同样大小的空图片
    mask=np.zeros_like(frame)
    # 填充多边形
    cv2.fillPoly(mask, pts=[np.array(polygonPoints)], color=(0, 255, 255))
    # 融合两张图片
    frame=cv2.addWeighted(frame, 0.7, mask, 0.3, 0)

    tracks=model.track(frame,persist=True,stream=True)
    frame=counter.count(frame)

    video_write.write(frame)
    cv2.imshow(winname='window',mat=frame)

    # 参数如果设置得太小,视频就会播放得很快
    if cv2.waitKey(1) & 0xFF == ord(' '):
        break

print('视频播放结束')

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