Support Inference over batch with TensorRT Engine Model

rtmo_gpu.py

@@ -242,7 +242,6 @@ def is_onnx_model(model_path):
         ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
         return True
     except Exception as e:
-        print('Error:', type(e))
         return False
 
 def is_trt_engine(model_path):
@@ -513,20 +512,35 @@ class RTMO_GPU_Batch(RTMO_GPU):
 
         input = batch_img
 
-        # Create an IO Binding object
-        io_binding = self.session.io_binding()
+        if self.model_format == 'onnx':
 
-        # Bind the model inputs and outputs to the IO Binding object
-        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')
+            # 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 postprocess_batch(