from collections import defaultdict
import cv2
from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator, colors

track_history = defaultdict(lambda: [])

model = YOLO("./Model/yolov8n-seg.pt")
VIDEO_PATH='./Test_Video/traffic_video.mp4'
RESULT_PATH='Object_tracking.mp4'
capture= cv2.VideoCapture(VIDEO_PATH)

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

video_writer=cv2.VideoWriter(RESULT_PATH, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))

while True:
    success, im0 =capture.read()
    if not success:
        print('视频读取完成')
        break

    annotator = Annotator(im0, line_width=2)
    results = model.track(im0, persist=True)

    if results[0].boxes.id is not None and results[0].masks is not None:
        masks = results[0].masks.xy
        track_ids = results[0].boxes.id.int().cpu().tolist()

        for mask,track_id in zip(masks, track_ids):
            color = colors(int(track_id), True)
            txt_color = annotator.get_txt_color(color)
            annotator.seg_bbox(mask=mask, mask_color=color, label=str(track_id), txt_color=txt_color)

    video_writer.write(im0)
    cv2.imshow(winname='windows',mat=im0)

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

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