tracking yolo

gistfile1.txt

from collections import defaultdict

import cv2

from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator, colors

track_history = defaultdict(lambda: [])

model = YOLO("/home/schoch/rocm_pytorch/runs/segment/train28/weights/best.pt")  # segmentation model
#cap = cv2.VideoCapture("/home/schoch/rocm_pytorch/datasets/v.mp4")
cap = cv2.VideoCapture("/mnt/j/yolo_pose_images/swing1_rl.mp4")
#cap = cv2.VideoCapture( "/mnt/c/Users/jesse/Documents/Kinovea/2024/202410/20241010-left/20241010-141751-left.mp4")
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))

out = cv2.VideoWriter("/mnt/j/yolo_pose_images/instance-segmentation-object-tracking6.avi", cv2.VideoWriter_fourcc(*"MJPG"), fps, (w, h))

while True:
    ret, im0 = cap.read()
    if not ret:
        print("Video frame is empty or video processing has been successfully completed.")
        break

    annotator = Annotator(im0, line_width=2)

    #results = model.track(im0, persist=True, tracker="bytetrack.yaml", conf=0.2)
    results = model.track(im0, persist=True, tracker="botsort.yaml", conf=0.1)

    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)

    out.write(im0)
    #cv2.imshow("instance-segmentation-object-tracking", im0)

    #if cv2.waitKey(1) & 0xFF == ord("q"):
        #break

out.release()
cap.release()