from ultralytics import YOLO
from ultralytics.solutions import object_counter
import cv2

model = YOLO("/home/daitranskku/code/cvpr2024/aicity/ultralytics/yolov8x.pt")
cap = cv2.VideoCapture("/home/daitranskku/code/SmartInsideAI/projects/STEAM/sample_videos/1699918175350_0.MOV")
assert cap.isOpened(), "Error reading video file"
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))

line_points = [(400, 100), (400, 1920)]  # line or region points
classes_to_count = [2, 7]  # person and car classes for count

# Video writer
video_writer = cv2.VideoWriter("object_counting_output.avi",
                       cv2.VideoWriter_fourcc(*'mp4v'),
                       fps,
                       (w, h))

# Init Object Counter
counter = object_counter.ObjectCounter()
counter.set_args(view_img=True,
                 reg_pts=line_points,
                 classes_names=model.names,
                 draw_tracks=True,
                 line_thickness=3)

while cap.isOpened():
    success, im0 = cap.read()
    # draw line
    cv2.line(im0, line_points[0], line_points[1], (0, 0, 255), 2)
    if not success:
        print("Video frame is empty or video processing has been successfully completed.")
        break
    tracks = model.track(im0, persist=True, show=True,
                         classes=classes_to_count)

    im0 = counter.start_counting(im0, tracks)
    video_writer.write(im0)

cap.release()
video_writer.release()
cv2.destroyAllWindows()