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

# Load a pretrained YOLO11n model
model = YOLO("yolo11m-2_uniform.onnx")
model_pose = YOLO("yolo11s-pose.pt")

video_path = "PF-071124-2.mp4"
cap = cv2.VideoCapture(video_path)

ret = True
frame_count = 0  # Initialize a frame counter
save_dir = "/results/"  # Directory to save the crops
os.makedirs(save_dir, exist_ok=True)  # Create the directory if it doesn't exist

while ret:
    ret, frame = cap.read()
    frame_count += 1  # Increment the frame counter
    #frame2_, frame_ = np.empty(2)
    if ret:
        results = model.predict(frame, conf=0.5)
        frame_ = results[0].plot()
        pose_results = model_pose.predict(frame_, show_boxes=False, show_conf=True, show_labels=True)
        results[0].keypoints = pose_results[0].keypoints

        frame_ = results[0].plot()

        
        cv2.imshow("frame", frame_)
            
        if cv2.waitKey(25) & 0xFF == ord('q'):
            break

cap.release()
cv2.destroyAllWindows()