import cv2
from ultralytics import YOLO

# Load models
det_model = YOLO('yolov8x.pt')
pose_model = YOLO('yolov8x-pose.pt')

input_video_path = 'cctv1_꽁초.mp4'
output_video_path = 'out_'+input_video_path

# Open input video
cap = cv2.VideoCapture(input_video_path)

# Get video properties
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(cap.get(cv2.CAP_PROP_FPS))

# Define the codec and create a VideoWriter object
out = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'MP4V'), fps, (width, height))

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # Run inference
    det_results = det_model(frame)
    pose_results = pose_model(frame)

    for result in det_results + pose_results:
        boxes = result.boxes
        keypoints = result.keypoints

        print(boxes)
        print(keypoints)

        print(boxes.data)
        print(keypoints.xywh)
        
        for box in boxes:
            cv2.rectangle(frame, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (255, 0, 0), 2)
    
        for keypoint in keypoints:
            for point in keypoint:
                cv2.circle(frame, (int(point[0]), int(point[1])), 3, (0, 0, 255), -1)
    
    frame_with_results = frame


    # Write the frame into the output video
    out.write(frame_with_results)

# Release video objects
cap.release()
out.release()
cv2.destroyAllWindows()