Add support TensorRT engine support for RTMO

rtmo_demo.py

@@ -12,18 +12,18 @@ 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)')
+    parser.add_argument('model_path', type=str, help='Path to a RTMO ONNX (or engine) model file (required)')
     parser.add_argument('--yolo_nas_pose', action='store_true', help='Use YOLO NAS Pose (flat format only) instead of RTMO Model')
 
     # Parse the command-line arguments
     args = parser.parse_args()
 
-    onnx_model = args.model_path # 'rtmo-s_8xb32-600e_body7-640x640.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() and not args.yolo_nas_pose else (640,640)
+    model_input_size = (416,416) if 'rtmo-t' in model.lower() and not args.yolo_nas_pose else (640,640)
 
-    body = RTMO_GPU(onnx_model=onnx_model, 
+    body = RTMO_GPU(model=model, 
                     model_input_size=model_input_size, is_yolo_nas_pose=args.yolo_nas_pose)
 
     for mp4_path in Path(args.path).glob('*'):
@@ -55,5 +55,5 @@ if __name__ == "__main__":
                                     kpt_thr=0.3,
                                     line_width=2)
             img_show = cv2.resize(img_show, (788, 525))
-            cv2.imshow(f'{onnx_model}', img_show)
+            cv2.imshow(f'{model}', img_show)
             cv2.waitKey(10)

rtmo_gpu.py

@@ -334,32 +334,48 @@ class RTMO_GPU(object):
             Returns:
                 outputs (np.ndarray): Output of RTMPose model.
             """
+
             # build input to (1, 3, H, W)
             img = img.transpose(2, 0, 1)
             img = np.ascontiguousarray(img, dtype=np.float32 if not self.is_yolo_nas_pose else np.uint8)
             input = img[None, :, :, :]
 
-            # Create an IO Binding object
-            io_binding = self.session.io_binding()
+            if self.model_format == 'onnx':
 
-            if not self.is_yolo_nas_pose:
-                # RTMO
-                io_binding.bind_input(name='input', device_type='cpu', device_id=0, element_type=np.float32, shape=input.shape, buffer_ptr=input.ctypes.data)
-                io_binding.bind_output(name='dets')
-                io_binding.bind_output(name='keypoints')
-            else:
-                # NAS Pose, flat format
-                io_binding.bind_input(name='input', device_type='cpu', device_id=0, element_type=np.uint8, shape=input.shape, buffer_ptr=input.ctypes.data)
-                io_binding.bind_output(name='graph2_flat_predictions')
+                # Create an IO Binding object
+                io_binding = self.session.io_binding()
 
-            # Run inference with IO Binding
-            self.session.run_with_iobinding(io_binding)
+                if not self.is_yolo_nas_pose:
+                    # RTMO
+                    io_binding.bind_input(name='input', device_type='cpu', device_id=0, element_type=np.float32, shape=input.shape, buffer_ptr=input.ctypes.data)
+                    io_binding.bind_output(name='dets')
+                    io_binding.bind_output(name='keypoints')
+                else:
+                    # NAS Pose, flat format
+                    io_binding.bind_input(name='input', device_type='cpu', device_id=0, element_type=np.uint8, shape=input.shape, buffer_ptr=input.ctypes.data)
+                    io_binding.bind_output(name='graph2_flat_predictions')
+
+                # Run inference with IO Binding
+                self.session.run_with_iobinding(io_binding)
+
+                # Retrieve the outputs from the IO Binding object
+                outputs = [output.numpy() for output in io_binding.get_outputs()]
+
+            else: # 'engine'
+
+                if not self.session.is_active:
+                    self.session.activate()
+
+                outputs = self.session.infer(feed_dict={'input': input}, check_inputs=False)
+                outputs = [output for output in outputs.values()]
 
-            # Retrieve the outputs from the IO Binding object
-            outputs = [output.numpy() for output in io_binding.get_outputs()]
-        
             return outputs
 
+    def __exit__(self):
+        if self.model_format == 'engine':
+            if self.session.is_active:
+                self.session.deactivate()
+
     def __call__(self, image: np.ndarray):
             image, ratio = self.preprocess(image)
 
@@ -371,33 +387,50 @@ class RTMO_GPU(object):
             return keypoints, scores
     
     def __init__(self,
-                 onnx_model: str = None,
+                 model: str = None,
                  model_input_size: tuple = (640, 640),
                  mean: tuple = None,
                  std: tuple = None,
                  device: str = 'cuda',
                  is_yolo_nas_pose = False):
-        
-        if not os.path.exists(onnx_model):
+
+        if not os.path.exists(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
-                                }),
-                                'CPUExecutionProvider']}
-        
-        self.session = ort.InferenceSession(path_or_bytes=onnx_model,
-                                            providers=providers[device])
+            raise FileNotFoundError(f"The specified ONNX model file was not found: {model}")
+
+        self.model = model
+        if model.endswith('.onnx'):
+            self.model_format = 'onnx'
+        elif model.endswith('.engine'):
+            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'])
+        else:
+            raise TypeError("Your model is neither ONNX nor Engine !")
+
+
+        if self.model_format == 'onnx':
+
+            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
+                                  }),
+                                  'CPUExecutionProvider']}
+
+            self.session = ort.InferenceSession(path_or_bytes=model,
+                                                providers=providers[device])
+
+        else: # 'engine'
+            engine = EngineFromBytes(BytesFromPath(model))
+            self.session = TrtRunner(engine)
 
-        self.onnx_model = onnx_model
         self.model_input_size = model_input_size
         self.mean = mean
         self.std = std