Add 'get_model_format_and_input_shape' helper function

rtmo_gpu.py

@@ -299,6 +299,19 @@ def get_trt_input_shapes(model_path):
                 input_shapes[binding] = engine.get_binding_shape(binding)
     return input_shapes
 
+def get_model_format_and_input_shape(model):
+    if is_onnx_model(model):
+        model_format = 'onnx'
+        input_shape = get_onnx_input_shapes(model)['input']
+    elif is_trt_engine(model):
+        model_format = 'engine'
+        from polygraphy.backend.trt import load_plugins
+        load_plugins(plugins=['libmmdeploy_tensorrt_ops.so'])
+        input_shape = get_trt_input_shapes(model)['input']
+    else:
+        raise TypeError("Your model is neither ONNX nor Engine !")
+    return model_format, input_shape
+
 class RTMO_GPU(object):
 
     def preprocess(self, img: np.ndarray):
@@ -465,18 +478,7 @@ class RTMO_GPU(object):
             raise FileNotFoundError(f"The specified ONNX model file was not found: {model}")
 
         self.model = model
-        if is_onnx_model(model):
-            self.model_format = 'onnx'
-            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 !")
-
+        self.model_format, self.input_shape = get_model_format_and_input_shape(self.model)
 
         if self.model_format == 'onnx':
 
@@ -497,6 +499,8 @@ class RTMO_GPU(object):
                                                 providers=providers[device])
 
         else: # 'engine'
+            from polygraphy.backend.common import BytesFromPath
+            from polygraphy.backend.trt import EngineFromBytes, TrtRunner
             engine = EngineFromBytes(BytesFromPath(model))
             self.session = TrtRunner(engine)