Repair DEMO to run with TRT EP

demo.sh

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

demo_batch.sh

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

fix_batch_dimension.py

@@ -0,0 +1,35 @@
+import argparse
+import onnx
+from onnx import numpy_helper
+import numpy as np
+
+def fix_batch_dimension(input_model_path, output_model_path, batch_size=1):
+    # Load the input ONNX model
+    model = onnx.load(input_model_path)
+
+    # Iterate through the model's inputs
+    for input_tensor in model.graph.input:
+        # Get the tensor shape
+        tensor_shape = input_tensor.type.tensor_type.shape
+        
+        # Modify the batch dimension
+        if len(tensor_shape.dim) > 0:
+            tensor_shape.dim[0].dim_value = batch_size
+
+    # Save the modified model to the output path
+    onnx.save(model, output_model_path)
+    print(f"Model saved with updated batch size of {batch_size} to {output_model_path}")
+
+if __name__ == "__main__":
+    # Setup command line arguments
+    parser = argparse.ArgumentParser(description="Fix batch dimension of an ONNX model.")
+    parser.add_argument("input_model_path", type=str, help="Path to the input ONNX model.")
+    parser.add_argument("output_model_path", type=str, help="Path to save the output ONNX model with fixed batch dimension.")
+    parser.add_argument("--batch_size", type=int, default=1, help="Value of batch size to assign (default is 1).")
+
+    # Parse arguments
+    args = parser.parse_args()
+
+    # Call the function to fix the batch dimension
+    fix_batch_dimension(args.input_model_path, args.output_model_path, args.batch_size)
+

fix_batch_dimension_all.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+# Loop through all .onnx files in the current directory
+for model in ./*.onnx; do
+    # Exclude processing any files that already have a batch size in their name
+    if [[ ! $model =~ _batch[0-9]+\.onnx$ ]]; then
+        # Process for batch size 1
+        python fix_batch_dimension.py "$model" "${model%.onnx}_batch1.onnx" --batch_size 1
+        echo "Generated ${model%.onnx}_batch1.onnx"
+
+        # Process for batch size 2
+        python fix_batch_dimension.py "$model" "${model%.onnx}_batch2.onnx" --batch_size 2
+        echo "Generated ${model%.onnx}_batch2.onnx"
+
+        # Process for batch size 4
+        python fix_batch_dimension.py "$model" "${model%.onnx}_batch4.onnx" --batch_size 4
+        echo "Generated ${model%.onnx}_batch4.onnx"
+    fi
+done
+

rtmo-l.fp16_batch1.onnx

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

rtmo-l.fp16_batch2.onnx

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

rtmo-l.fp16_batch4.onnx

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

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:cac7daf35257f78c533e088773df544806c69d03f12aef9937ecf0e161649e4a
+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_batch1.onnx

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

rtmo-m.fp16_batch2.onnx

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

rtmo-m.fp16_batch4.onnx

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

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:f64f9494e826b39b8347af6de0e9fcb932c58762578fdcf4cdcc53be5852bacf
+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_batch1.onnx

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

rtmo-s.fp16_batch2.onnx

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

rtmo-s.fp16_batch4.onnx

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

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:28a5f6e563ad3a20c9e2c18355514c0a2c0aa147f09017f2fbace79a9d7ba38d
+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_batch1.onnx

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

rtmo-t.fp16_batch2.onnx

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

rtmo-t.fp16_batch4.onnx

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

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:21d77708f79cd81b644f53804f6541d672e10556d789cec8eebbeafc9d41a8ce
+size 27375829

rtmo-t_batch4.onnx

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