Make ONNX models compatible with ONNXruntime's TensorrtExecutionProvider

- Use FP16 ONNX Model with Batch Size 4 for DEMO with TRT EP
- Perform shape inference on all ONNX models
- Removed int8 models which is not recommended due to performance drop
- Add ONNX models with fixed input batch size 1,2,4
- Enable TensorRT Execution Provider (TRT EP)

demo.sh

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

demo_batch.sh

@@ -0,0 +1,2 @@
+#!/bin/sh
+python3 rtmo_demo_batch.py ./video rtmo-t.fp16_batch4.onnx 4

fix_input_batch_size.py

@@ -0,0 +1,31 @@
+import onnx
+import argparse
+
+def modify_batch_size(input_model_path, output_model_path, batch_size):
+    # Load the ONNX model
+    model = onnx.load(input_model_path)
+
+    # Modify the batch size of the first input tensor
+    input_tensor = model.graph.input[0]
+    input_tensor.type.tensor_type.shape.dim[0].dim_value = batch_size
+
+    # Save the modified model
+    onnx.save(model, output_model_path)
+    print(f"Modified model saved to {output_model_path}")
+
+if __name__ == "__main__":
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description="Modify the batch size of the first input tensor of an ONNX model.")
+    parser.add_argument("input_model_path", type=str, help="Path to the input ONNX model file.")
+    parser.add_argument("output_model_path", type=str, help="Path to save the modified ONNX model file.")
+    parser.add_argument("--batch_size", type=int, help="Desired batch size for the first input tensor.")
+
+    args = parser.parse_args()
+
+    # Ensure that the batch_size argument has been provided
+    if args.batch_size is None:
+        raise ValueError("Please specify the new batch size for the input tensor using --batch_size.")
+
+    # Call the function to modify the model's batch size
+    modify_batch_size(args.input_model_path, args.output_model_path, args.batch_size)
+

fix_input_batch_size_all.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+
+# Directory containing the original ONNX models
+MODEL_DIR="."
+
+# Directory to save the modified ONNX models
+OUTPUT_DIR="."
+
+# Ensure output directory exists
+mkdir -p "$OUTPUT_DIR"
+
+# Path to the Python script
+PYTHON_SCRIPT="fix_input_batch_size.py"
+
+# Loop through each ONNX model in the directory
+for model in "$MODEL_DIR"/*.onnx; do
+    echo "Processing $model..."
+    filename=$(basename -- "$model")
+    base="${filename%.*}"
+    
+    # Generate models with batch size 1
+    output_model_1="$OUTPUT_DIR/${base}_batch1.onnx"
+    python "$PYTHON_SCRIPT" "$model" "$output_model_1" --batch_size 1
+    echo "Generated model with batch size 1: $output_model_1"
+
+    # Generate models with batch size 2
+    output_model_2="$OUTPUT_DIR/${base}_batch2.onnx"
+    python "$PYTHON_SCRIPT" "$model" "$output_model_2" --batch_size 1
+    echo "Generated model with batch size 2: $output_model_2"
+    
+    # Generate models with batch size 4
+    output_model_4="$OUTPUT_DIR/${base}_batch4.onnx"
+    python "$PYTHON_SCRIPT" "$model" "$output_model_4" --batch_size 4
+    echo "Generated model with batch size 4: $output_model_4"
+done
+
+echo "Batch size modification complete."
+

rtmo-l.fp16.onnx

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

rtmo-l.int8.onnx → rtmo-l.fp16_batch1.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:21421a36217dc3545201fab4d7fde993636eab0b46715d97cd101fda91fc834f
-size 44653310
+oid sha256:c9c9fee334086aec0482b348cb74a61bc144c0a9bd23550ebdc05001f2ca86d9
+size 88036970

rtmo-m.int8.onnx → rtmo-l.fp16_batch2.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:eba6bdd56b087e9111d17c5421a1a9896f064e64b60d7192a1c15cba6623b5fc
-size 22942317
+oid sha256:c9c9fee334086aec0482b348cb74a61bc144c0a9bd23550ebdc05001f2ca86d9
+size 88036970

rtmo-s.int8.onnx → rtmo-l.fp16_batch4.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:02e26bfe3cc468e9e49942eee10f62d2435bdd4c6ad92b12aafb7e2deae85ddc
-size 10474466
+oid sha256:0f2098f33be14598dee33f0203a6fa3ed72e80159b691e1c664ef9869a7ec215
+size 88036970

rtmo-l.onnx

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

rtmo-l_batch1.onnx

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

rtmo-l_batch2.onnx

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

rtmo-l_batch4.onnx

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

rtmo-m.fp16.onnx

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

rtmo-t.int8.onnx → rtmo-m.fp16_batch1.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2dc34c21f47fddd017f31dd8b0a375a2754eae06b0b29732c5a6e740c45f9cd5
-size 7210158
+oid sha256:59d4a54c134dabeecbe58597b323be6e31ae500745476f2f276e38a2b03c4ed9
+size 44714757

rtmo-m.fp16_batch2.onnx

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

rtmo-m.fp16_batch4.onnx

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

rtmo-m.onnx

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

rtmo-m_batch1.onnx

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

rtmo-m_batch2.onnx

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

rtmo-m_batch4.onnx

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

rtmo-s.fp16.onnx

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

rtmo-s.fp16_batch1.onnx

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

rtmo-s.fp16_batch2.onnx

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

rtmo-s.fp16_batch4.onnx

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

rtmo-s.onnx

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

rtmo-s_batch1.onnx

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

rtmo-s_batch2.onnx

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

rtmo-s_batch4.onnx

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

rtmo-t.fp16.onnx

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

rtmo-t.fp16_batch1.onnx

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

rtmo-t.fp16_batch2.onnx

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

rtmo-t.fp16_batch4.onnx

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

rtmo-t.onnx

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

rtmo-t_batch1.onnx

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

rtmo-t_batch2.onnx

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

rtmo-t_batch4.onnx

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

rtmo_demo_batch.py

@@ -6,12 +6,11 @@ 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):
+def process_video(video_path, body_estimator, batch_size=4):
     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()
@@ -61,6 +60,7 @@ if __name__ == "__main__":
     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()
@@ -72,4 +72,4 @@ if __name__ == "__main__":
     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)
+        process_video(str(mp4_path), body_estimator, args.batch_size)

rtmo_gpu.py

@@ -358,11 +358,15 @@ class RTMO_GPU(object):
                  device: str = 'cuda'):
         
         if not os.path.exists(onnx_model):
-            from rtmlib.tools.file import download_checkpoint
-            onnx_model = download_checkpoint(onnx_model)
+            # If the file does not exist, raise FileNotFoundError
+            raise FileNotFoundError(f"The specified ONNX model file was not found: {onnx_model}")
 
         providers = {'cpu': 'CPUExecutionProvider',
                         'cuda': [
+                                    ('TensorrtExecutionProvider', {
+                                    'trt_fp16_enable':True,
+                                    'trt_engine_cache_enable':True,
+                                    'trt_engine_cache_path':'cache'}),
                                     ('CUDAExecutionProvider', {
                                     'cudnn_conv_algo_search': 'DEFAULT',
                                     'cudnn_conv_use_max_workspace': True

symbolic_shape_infer_all.sh

@@ -0,0 +1,6 @@
+for f in $(ls -S *.onnx | tac); 
+do 
+echo Process "$f"
+python3 symbolic_shape_infer.py --input "$f" --output "$f" --auto_merge; 
+done
+