Add model format (ONNX/Engine) & input size detection (based on file header, not on filename) for RTMO

rtmo_demo.py

@@ -20,11 +20,7 @@ if __name__ == "__main__":
 
     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 model.lower() and not args.yolo_nas_pose else (640,640)
-
-    body = RTMO_GPU(model=model, 
-                    model_input_size=model_input_size, is_yolo_nas_pose=args.yolo_nas_pose)
+    body = RTMO_GPU(model=model, is_yolo_nas_pose=args.yolo_nas_pose)
 
     for mp4_path in Path(args.path).glob('*'):
     

rtmo_demo_batch.py

@@ -74,10 +74,9 @@ if __name__ == "__main__":
     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)
+    body_estimator = RTMO_GPU_Batch(model=onnx_model)
 
     for mp4_path in Path(args.path).glob('*'):
         process_video(str(mp4_path), body_estimator, args.batch_size)

rtmo_gpu.py

@@ -237,6 +237,43 @@ def draw_skeleton(img,
         raise NotImplementedError
     return img
 
+def is_onnx_model(model_path):
+    try:
+        import onnx
+        onnx_model = onnx.load(model_path)
+        onnx.checker.check_model(onnx_model)
+        return True
+    except Exception as e:
+        return False
+
+def is_trt_engine(model_path):
+    try:
+        from polygraphy.backend.common import BytesFromPath
+        from polygraphy.backend.trt import EngineFromBytes
+        engine = EngineFromBytes(BytesFromPath(model_path))
+        return engine is not None
+    except Exception:
+        return False
+
+def get_onnx_input_shapes(model_path):
+    from polygraphy.backend.onnx.loader import OnnxFromPath
+    from polygraphy.backend.onnx import infer_shapes
+    model = OnnxFromPath(model_path)()
+    model = infer_shapes(model)
+    input_shapes = {inp.name: inp.type.tensor_type.shape for inp in model.graph.input}
+    return {name: [dim.dim_value if dim.dim_value > 0 else 'Dynamic' for dim in shape_proto.dim] 
+            for name, shape_proto in input_shapes.items()}
+
+def get_trt_input_shapes(model_path):
+    input_shapes = {}
+    import tensorrt as trt
+    with open(model_path, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.WARNING)) as runtime:
+        engine = runtime.deserialize_cuda_engine(f.read())
+        for binding in engine:
+            if engine.binding_is_input(binding):
+                input_shapes[binding] = engine.get_binding_shape(binding)
+    return input_shapes
+
 class RTMO_GPU(object):
 
     def preprocess(self, img: np.ndarray):
@@ -388,7 +425,6 @@ class RTMO_GPU(object):
     
     def __init__(self,
                  model: str = None,
-                 model_input_size: tuple = (640, 640),
                  mean: tuple = None,
                  std: tuple = None,
                  device: str = 'cuda',
@@ -399,13 +435,15 @@ class RTMO_GPU(object):
             raise FileNotFoundError(f"The specified ONNX model file was not found: {model}")
 
         self.model = model
-        if model.endswith('.onnx'):
+        if is_onnx_model(model):
             self.model_format = 'onnx'
-        elif model.endswith('.engine'):
+            self.input_shape = get_onnx_input_shapes(self.model)['input']
+        elif is_trt_engine(model):
             self.model_format = 'engine'
             from polygraphy.backend.common import BytesFromPath
             from polygraphy.backend.trt import EngineFromBytes, TrtRunner, load_plugins
             load_plugins(plugins=['libmmdeploy_tensorrt_ops.so'])
+            self.input_shape = get_trt_input_shapes(self.model)['input']
         else:
             raise TypeError("Your model is neither ONNX nor Engine !")
 
@@ -431,7 +469,7 @@ class RTMO_GPU(object):
             engine = EngineFromBytes(BytesFromPath(model))
             self.session = TrtRunner(engine)
 
-        self.model_input_size = model_input_size
+        self.model_input_size = self.input_shape[2:4] # B, C, H, W,
         self.mean = mean
         self.std = std
         self.device = device