| #!/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 (or engine) 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() | |
| model = args.model_path # 'rtmo-s_8xb32-600e_body7-640x640.onnx' | |
| body = RTMO_GPU(model=model, 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'{model}', img_show) | |
| cv2.waitKey(10) | |