#!/usr/bin/python3

import time
import cv2
from pathlib import Path
import argparse
import os
from rtmo_gpu import RTMO_GPU, draw_skeleton

if __name__ == "__main__":

    # Set up argument parsing
    parser = argparse.ArgumentParser(description='Process the path to a video file folder.')
    parser.add_argument('path', type=str, help='Path to the folder containing video files (required)')
    parser.add_argument('model_path', type=str, help='Path to a RTMO ONNX model file (required)')
    parser.add_argument('--yolo_nas_pose', action='store_true', help='Use YOLO NAS Pose (flat format only) instead of RTMO Model')

    # Parse the command-line arguments
    args = parser.parse_args()

    onnx_model = args.model_path # 'rtmo-s_8xb32-600e_body7-640x640.onnx'

    # Only Tiny Model has (416,416) as input model
    model_input_size = (416,416) if 'rtmo-t' in onnx_model.lower() and not args.yolo_nas_pose else (640,640)

    body = RTMO_GPU(onnx_model=onnx_model, 
                    model_input_size=model_input_size, is_yolo_nas_pose=args.yolo_nas_pose)

    for mp4_path in Path(args.path).glob('*'):
    
        # Now, use the best.url, which is the direct video link for streaming
        cap = cv2.VideoCapture(filename=os.path.abspath(mp4_path))

        frame_idx = 0

        while cap.isOpened():
            success, frame = cap.read()
            frame_idx += 1

            if not success:
                break
            s = time.time()
            keypoints, scores = body(frame)
            det_time = time.time() - s
            print(f'det: {round(1.0 / det_time,1)} FPS')
            
            img_show = frame.copy()
            
            # if you want to use black background instead of original image,
            # img_show = np.zeros(img_show.shape, dtype=np.uint8)

            img_show = draw_skeleton(img_show,
                                    keypoints,
                                    scores,
                                    kpt_thr=0.3,
                                    line_width=2)
            img_show = cv2.resize(img_show, (788, 525))
            cv2.imshow(f'{onnx_model}', img_show)
            cv2.waitKey(10)