RTMO ONNX Models from MMPOSE range from Tiny to Large

convert_to_fp16.py

@@ -0,0 +1,30 @@
+import argparse
+from onnxmltools.utils.float16_converter import convert_float_to_float16
+from onnxmltools.utils import load_model, save_model
+
+def main():
+    # Set up an argument parser
+    parser = argparse.ArgumentParser(description='Convert ONNX model from Float32 to Float16.')
+    parser.add_argument('--input_model', type=str, required=True, help='Path to the input ONNX model file.')
+    parser.add_argument('--output_model', type=str, required=True, help='Path for saving the converted ONNX model file.')
+    
+    # Parse arguments
+    args = parser.parse_args()
+
+    # Load the model
+    print(f"Loading model from {args.input_model}")
+    onnx_model = load_model(args.input_model)
+
+    # Convert model from Float32 to Float16
+    print("Converting model...")
+    new_onnx_model = convert_float_to_float16(onnx_model, keep_io_types=True)
+
+    # Save the converted model
+    print(f"Saving converted model to {args.output_model}")
+    save_model(new_onnx_model, args.output_model)
+
+    print("Conversion complete.")
+
+if __name__ == "__main__":
+    main()
+

demo.sh

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

rtmo-l.fp16.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7252d22c1aa13abbd2b8616fb8468ba90465714719da9f8821c7fa25e6080aa2
+size 88026530

rtmo-l.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:090096ca90f29163cc4f67137dcc0cd4b2ee95ea0af11764fbfda88dd2ae1140
+size 175901910

rtmo-m.fp16.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9bcaef995703f899258c04fd12a5bdd773e7f783e60d4b10b4a1a744caa1557
+size 44704317

rtmo-m.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:750d872151ff4652f02c738efc3a547e112ce9b688fe920bc17f948e8c3afdac
+size 89269977

rtmo-s.fp16.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eda631cea4ce08c647dd76b835b290413191567acb0b2663c8e980318ef5d427
+size 19871080

rtmo-s.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d0703d40d19f3921da51ae725402d5fdae4d2478c7442072d3101bd396f370d8
+size 39617685

rtmo-t.fp16.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c45e15eb6609e4893697e5bf5ef2d40c8b6e15c77dbb30e72f148c2d0ca46ea
+size 13733099

rtmo-t.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:20aad6e2e42359cac1c5b4a0b2da00e29bfe91a72a782fdcf287d273a04c1b24
+size 27342008

rtmo_demo.py

@@ -0,0 +1,60 @@
+#!/usr/bin/python3
+
+import time
+import cv2
+from rtmlib import draw_skeleton
+from pathlib import Path
+import argparse
+import os
+from rtmo_gpu import RTMO_GPU
+
+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 # '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() else (640,640)
+
+    body = RTMO_GPU(onnx_model=onnx_model, 
+                    model_input_size=model_input_size)
+
+    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,
+                                    openpose_skeleton=False,
+                                    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)

rtmo_gpu.py

@@ -0,0 +1,55 @@
+import os
+from rtmlib import RTMO
+
+class RTMO_GPU(RTMO):
+
+    def __init__(self,
+                 onnx_model: str = None,
+                 model_input_size: tuple = (640, 640),
+                 mean: tuple = None,
+                 std: tuple = None,
+                 to_openpose: bool = False,
+                 backend: str = 'onnxruntime',
+                 device: str = 'cuda'):
+        
+        if backend == 'onnxruntime':
+
+            if not os.path.exists(onnx_model):
+                from rtmlib.tools.file import download_checkpoint
+                onnx_model = download_checkpoint(onnx_model)
+
+            import onnxruntime as ort
+
+            providers = {'cpu': 'CPUExecutionProvider',
+                         'cuda': [
+                                      ('CUDAExecutionProvider', {
+                                        'device_id': 0,
+                                        'arena_extend_strategy': 'kNextPowerOfTwo',
+                                        'gpu_mem_limit': 2 * 1024 * 1024 * 1024,
+                                        'cudnn_conv_algo_search': 'DEFAULT',
+                                        'do_copy_in_default_stream': True,
+                                        'enable_cuda_graph': False
+                                    }),
+                                  'CPUExecutionProvider']}
+            
+            self.session = ort.InferenceSession(path_or_bytes=onnx_model,
+                                                providers=providers[device])
+
+            print(f'load {onnx_model} with {backend} backend')
+
+            self.onnx_model = onnx_model
+            self.model_input_size = model_input_size
+            self.mean = mean
+            self.std = std
+            self.backend = backend
+            self.device = device
+            self.to_openpose = to_openpose
+
+        else:
+            super().__init__(onnx_model,
+                             model_input_size, 
+                             mean,
+                             std, 
+                             to_openpose,
+                             backend, 
+                             device)