File size: 3,240 Bytes
3939a91 35b2a45 3939a91 b75f05d 3939a91 42892de 3939a91 15801f5 3939a91 35b2a45 15801f5 3939a91 12e168d 42892de 3939a91 35b2a45 3939a91 b75f05d 3939a91 58a44cf 3939a91 b75f05d |
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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
#!/usr/bin/python3
import time
import cv2
from pathlib import Path
import argparse
from rtmo_gpu import RTMO_GPU_Batch, draw_skeleton, resize_to_fit_screen # Ensure to import RTMO_GPU_Batch
def process_video(video_path, body_estimator, batch_size=4):
cap = cv2.VideoCapture(video_path)
batch_frames = []
frame_idxs = []
while cap.isOpened():
success, frame = cap.read()
if not success:
break
batch_frames.append(frame)
frame_idxs.append(cap.get(cv2.CAP_PROP_POS_FRAMES))
# Process the batch when it's full
if len(batch_frames) == batch_size:
s = time.time()
batch_keypoints, batch_scores = body_estimator.__batch_call__(batch_frames)
det_time = time.time() - s
fps = round(batch_size / det_time, 1)
print(f'Batch det: {fps} FPS')
for i, keypoints in enumerate(batch_keypoints):
scores = batch_scores[i]
frame = batch_frames[i]
img_show = frame.copy()
img_show = draw_skeleton(img_show, keypoints, scores, kpt_thr=0.3, line_width=2)
img_show = resize_to_fit_screen(img_show, 720, 480)
cv2.putText(img_show, f'{fps:.1f}', (10, 30), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 255, 0), 1, cv2.LINE_AA)
cv2.imshow(f'{video_path}', img_show)
cv2.waitKey(10)
# Clear the batch
batch_frames = []
# Process remaining frames if any
if batch_frames:
# Padding
while len(batch_frames) < batch_size:
# Option 1: Add a black frame
# black_frame = np.zeros_like(batch_frames[0])
# batch_frames.append(black_frame)
# Option 2: Duplicate the last frame
batch_frames.append(batch_frames[-1])
batch_keypoints, batch_scores = body_estimator.__batch_call__(batch_frames)
for i, keypoints in enumerate(batch_keypoints):
scores = batch_scores[i]
frame = batch_frames[i]
img_show = frame.copy()
img_show = draw_skeleton(img_show, keypoints, scores, kpt_thr=0.3, line_width=2)
img_show = resize_to_fit_screen(img_show, 720, 480)
cv2.imshow(f'{video_path}', img_show)
#cv2.waitKey(10)
cap.release()
cv2.destroyAllWindows()
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('batch_size', type=int, help='Path to a RTMO ONNX input batch size (required)')
# Parse the command-line arguments
args = parser.parse_args()
onnx_model = args.model_path # Example: 'rtmo-s_8xb32-600e_body7-640x640.onnx'
# Instantiate the RTMO_GPU_Batch instead of RTMO_GPU
body_estimator = RTMO_GPU_Batch(model=onnx_model)
for mp4_path in Path(args.path).glob('*'):
process_video(str(mp4_path), body_estimator, args.batch_size)
|