File size: 3,148 Bytes
3939a91 b75f05d 3939a91 12e168d 3939a91 b75f05d 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 |
#!/usr/bin/python3
import time
import cv2
from pathlib import Path
import argparse
from rtmo_gpu import RTMO_GPU_Batch, draw_skeleton # 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_frames)
det_time = time.time() - s
print(f'Batch det: {round(batch_size / det_time, 1)} 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 = cv2.resize(img_show, (788, 525))
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_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 = cv2.resize(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'
model_input_size = (416, 416) if 'rtmo-t' in onnx_model.lower() else (640, 640)
# Instantiate the RTMO_GPU_Batch instead of RTMO_GPU
body_estimator = RTMO_GPU_Batch(onnx_model=onnx_model, model_input_size=model_input_size)
for mp4_path in Path(args.path).glob('*'):
process_video(str(mp4_path), body_estimator, args.batch_size)
|