Use RTMO_GPU_Batch to demo

demo.sh

@@ -1,2 +1,2 @@
 #!/bin/sh
-python3 rtmo_demo.py ./video rtmo-t.fp16.onnx
+python3 rtmo_demo_batch.py ./video rtmo-t.fp16.onnx

rtmo_demo_batch.py

@@ -0,0 +1,75 @@
+#!/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):
+    cap = cv2.VideoCapture(video_path)
+
+    batch_frames = []
+    frame_idxs = []
+    batch_size = 4  # Define a suitable batch size based on your GPU memory
+
+    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:
+        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)')
+
+    # 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)