File size: 2,182 Bytes
7ad7e4d 43a489d 7ad7e4d 1f0f5d8 7ad7e4d 1f0f5d8 7ad7e4d 1f0f5d8 7ad7e4d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
#!/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)
|