DeepStream 7.1 + Fixes + New model output format

utils/export_ppyoloe.py

@@ -1,14 +1,14 @@
 import os
-import sys
 import onnx
 import paddle
 import paddle.nn as nn
-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.slim import build_slim_model
 from ppdet.data.source.category import get_categories
+from ppdet.utils.check import check_version, check_config
+from ppdet.core.workspace import load_config, merge_config
 
 
 class DeepStreamOutput(nn.Layer):
@@ -18,9 +18,20 @@ class DeepStreamOutput(nn.Layer):
     def forward(self, x):
         boxes = x['bbox']
         x['bbox_num'] = x['bbox_num'].transpose([0, 2, 1])
-        scores = paddle.max(x['bbox_num'], 2, keepdim=True)
-        classes = paddle.cast(paddle.argmax(x['bbox_num'], 2, keepdim=True), dtype='float32')
-        return boxes, scores, classes
+        scores = paddle.max(x['bbox_num'], axis=-1, keepdim=True)
+        labels = paddle.argmax(x['bbox_num'], axis=-1, keepdim=True)
+        return paddle.concat((boxes, scores, paddle.cast(labels, dtype=boxes.dtype)), axis=-1)
+
+
+class DeepStreamInput(nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        y = {}
+        y['image'] = x['image']
+        y['scale_factor'] = paddle.to_tensor([1.0, 1.0], dtype=x['image'].dtype)
+        return y
 
 
 def ppyoloe_export(FLAGS):
@@ -43,10 +54,17 @@ def ppyoloe_export(FLAGS):
     return trainer.cfg, static_model
 
 
-def main(FLAGS):
-    print('\nStarting: %s' % FLAGS.weights)
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore')
 
-    print('\nOpening PPYOLOE model\n')
+
+def main(FLAGS):
+    suppress_warnings()
+
+    print(f'\nStarting: {FLAGS.weights}')
+
+    print('Opening PPYOLOE model')
 
     paddle.set_device('cpu')
     cfg, model = ppyoloe_export(FLAGS)
@@ -54,32 +72,30 @@ def main(FLAGS):
     anno_file = cfg['TestDataset'].get_anno()
     if os.path.isfile(anno_file):
         _, catid2name = get_categories(cfg['metric'], anno_file, 'detection_arch')
-        print('\nCreating labels.txt file')
-        f = open('labels.txt', 'w')
-        for name in catid2name.values():
-            f.write(str(name) + '\n')
-        f.close()
+        print('Creating labels.txt file')
+        with open('labels.txt', 'w', encoding='utf-8') as f:
+            for name in catid2name.values():
+                f.write(f'{name}\n')
 
-    model = nn.Sequential(model, DeepStreamOutput())
+    model = nn.Sequential(DeepStreamInput(), model, DeepStreamOutput())
 
     img_size = [cfg.eval_height, cfg.eval_width]
 
     onnx_input_im = {}
-    onnx_input_im['image'] = paddle.static.InputSpec(shape=[FLAGS.batch, 3, *img_size], dtype='float32', name='image')
-    onnx_input_im['scale_factor'] = paddle.static.InputSpec(shape=[FLAGS.batch, 2], dtype='float32', name='scale_factor')
-    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')
+    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():
@@ -92,9 +108,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
@@ -102,4 +118,4 @@ def parse_args():
 
 if __name__ == '__main__':
     FLAGS = parse_args()
-    sys.exit(main(FLAGS))
+    main(FLAGS)