DeepStream 7.1 + Fixes + New model output format

utils/export_rtdetr_paddle.py

@@ -1,34 +1,32 @@
 import os
-import sys
-import warnings
 import onnx
 import paddle
 import paddle.nn as nn
 import paddle.nn.functional as F
-from ppdet.core.workspace import load_config, merge_config
-from ppdet.utils.check import check_version, check_config
-from ppdet.utils.cli import ArgsParser
+
 from ppdet.engine import Trainer
+from ppdet.utils.cli import ArgsParser
+from ppdet.utils.check import check_version, check_config
+from ppdet.core.workspace import load_config, merge_config
 
 
 class DeepStreamOutput(nn.Layer):
     def __init__(self, img_size, use_focal_loss):
+        super().__init__()
         self.img_size = img_size
         self.use_focal_loss = use_focal_loss
-        super().__init__()
 
     def forward(self, x):
         boxes = x['bbox']
-        out_shape = paddle.to_tensor([[*self.img_size]]).flip(1).tile([1, 2]).unsqueeze(1)
-        boxes *= out_shape
+        convert_matrix = paddle.to_tensor(
+            [[1, 0, 1, 0], [0, 1, 0, 1], [-0.5, 0, 0.5, 0], [0, -0.5, 0, 0.5]], dtype=boxes.dtype
+        )
+        boxes @= convert_matrix
+        boxes *= paddle.to_tensor([[*self.img_size]]).flip(1).tile([1, 2]).unsqueeze(1)
         bbox_num = F.sigmoid(x['bbox_num']) if self.use_focal_loss else F.softmax(x['bbox_num'])[:, :, :-1]
-        scores = paddle.max(bbox_num, 2, keepdim=True)
-        classes = paddle.cast(paddle.argmax(bbox_num, 2, keepdim=True), dtype='float32')
-        return boxes, scores, classes
-    
-
-def suppress_warnings():
-    warnings.filterwarnings('ignore')
+        scores = paddle.max(bbox_num, axis=-1, keepdim=True)
+        labels = paddle.argmax(bbox_num, axis=-1, keepdim=True)
+        return paddle.concat((boxes, scores, paddle.cast(labels, dtype=boxes.dtype)), axis=-1)
 
 
 def rtdetr_paddle_export(FLAGS):
@@ -50,12 +48,17 @@ def rtdetr_paddle_export(FLAGS):
     return trainer.cfg, static_model
 
 
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore')
+
+
 def main(FLAGS):
     suppress_warnings()
 
-    print('\nStarting: %s' % FLAGS.weights)
+    print(f'\nStarting: {FLAGS.weights}')
 
-    print('\nOpening RT-DETR Paddle model\n')
+    print('Opening RT-DETR Paddle model')
 
     paddle.set_device('cpu')
     cfg, model = rtdetr_paddle_export(FLAGS)
@@ -65,20 +68,20 @@ def main(FLAGS):
     model = nn.Sequential(model, DeepStreamOutput(img_size, cfg.use_focal_loss))
 
     onnx_input_im = {}
-    onnx_input_im['image'] = paddle.static.InputSpec(shape=[FLAGS.batch, 3, *img_size], dtype='float32', name='image')
-    onnx_output_file = cfg.filename + '.onnx'
+    onnx_input_im['image'] = paddle.static.InputSpec(shape=[FLAGS.batch, 3, *img_size], dtype='float32')
+    onnx_output_file = f'{FLAGS.weights}.onnx'
 
-    print('\nExporting the model to ONNX\n')
-    paddle.onnx.export(model, cfg.filename, input_spec=[onnx_input_im], opset_version=FLAGS.opset)
+    print('Exporting the model to ONNX\n')
+    paddle.onnx.export(model, FLAGS.weights, input_spec=[onnx_input_im], opset_version=FLAGS.opset)
 
     if FLAGS.simplify:
-        print('\nSimplifying the ONNX model')
-        import onnxsim
+        print('Simplifying the ONNX model')
+        import onnxslim
         model_onnx = onnx.load(onnx_output_file)
-        model_onnx, _ = onnxsim.simplify(model_onnx)
+        model_onnx = onnxslim.slim(model_onnx)
         onnx.save(model_onnx, onnx_output_file)
 
-    print('\nDone: %s\n' % onnx_output_file)
+    print(f'Done: {onnx_output_file}\n')
 
 
 def parse_args():
@@ -91,9 +94,9 @@ def parse_args():
     parser.add_argument('--batch', type=int, default=1, help='Static batch-size')
     args = parser.parse_args()
     if not os.path.isfile(args.weights):
-        raise SystemExit('\nInvalid weights file')
+        raise SystemExit('Invalid weights file')
     if args.dynamic and args.batch > 1:
-        raise SystemExit('\nCannot set dynamic batch-size and static batch-size at same time')
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
     elif args.dynamic:
         args.batch = None
     return args
@@ -101,4 +104,4 @@ def parse_args():
 
 if __name__ == '__main__':
     FLAGS = parse_args()
-    sys.exit(main(FLAGS))
+    main(FLAGS)