DeepStream 7.1 + Fixes + New model output format

utils/export_damoyolo.py

@@ -1,14 +1,12 @@
 import os
-import sys
-import argparse
-import warnings
 import onnx
 import torch
 import torch.nn as nn
-from damo.base_models.core.ops import RepConv, SiLU
+
 from damo.config.base import parse_config
-from damo.detectors.detector import build_local_model
 from damo.utils.model_utils import replace_module
+from damo.base_models.core.ops import RepConv, SiLU
+from damo.detectors.detector import build_local_model
 
 
 class DeepStreamOutput(nn.Module):
@@ -17,15 +15,17 @@ class DeepStreamOutput(nn.Module):
 
     def forward(self, x):
         boxes = x[1]
-        scores, classes = torch.max(x[0], 2, keepdim=True)
-        classes = classes.float()
-        return boxes, scores, classes
+        scores, labels = torch.max(x[0], dim=-1, keepdim=True)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
 
 
 def suppress_warnings():
+    import warnings
     warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
     warnings.filterwarnings('ignore', category=UserWarning)
     warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
 
 
 def damoyolo_export(weights, config_file, device):
@@ -48,7 +48,7 @@ def damoyolo_export(weights, config_file, device):
 def main(args):
     suppress_warnings()
 
-    print('\nStarting: %s' % args.weights)
+    print(f'\nStarting: {args.weights}')
 
     print('Opening DAMO-YOLO model')
 
@@ -57,49 +57,44 @@ def main(args):
 
     if len(cfg.dataset['class_names']) > 0:
         print('Creating labels.txt file')
-        f = open('labels.txt', 'w')
-        for name in cfg.dataset['class_names']:
-            f.write(name + '\n')
-        f.close()
+        with open('labels.txt', 'w', encoding='utf-8') as f:
+            for name in cfg.dataset['class_names']:
+                f.write(f'{name}\n')
 
     model = nn.Sequential(model, DeepStreamOutput())
 
     img_size = args.size * 2 if len(args.size) == 1 else args.size
 
     onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
-    onnx_output_file = cfg.miscs['exp_name'] + '.onnx'
+    onnx_output_file = f'{args.weights}.onnx'
 
     dynamic_axes = {
         'input': {
             0: 'batch'
         },
-        'boxes': {
-            0: 'batch'
-        },
-        'scores': {
-            0: 'batch'
-        },
-        'classes': {
+        'output': {
             0: 'batch'
         }
     }
 
     print('Exporting the model to ONNX')
-    torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
-                      do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
-                      dynamic_axes=dynamic_axes if args.dynamic else None)
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
 
     if args.simplify:
         print('Simplifying the ONNX model')
-        import onnxsim
+        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('Done: %s\n' % onnx_output_file)
+    print(f'Done: {onnx_output_file}\n')
 
 
 def parse_args():
+    import argparse
     parser = argparse.ArgumentParser(description='DeepStream DAMO-YOLO conversion')
     parser.add_argument('-w', '--weights', required=True, help='Input weights (.pth) file path (required)')
     parser.add_argument('-c', '--config', required=True, help='Input config (.py) file path (required)')
@@ -120,4 +115,4 @@ def parse_args():
 
 if __name__ == '__main__':
     args = parse_args()
-    sys.exit(main(args))
+    main(args)

utils/export_goldyolo.py

@@ -0,0 +1,121 @@
+import os
+import onnx
+import torch
+import torch.nn as nn
+
+import yolov6.utils.general as _m
+from yolov6.layers.common import SiLU
+from gold_yolo.switch_tool import switch_to_deploy
+from yolov6.utils.checkpoint import load_checkpoint
+
+
+def _dist2bbox(distance, anchor_points, box_format='xyxy'):
+    lt, rb = torch.split(distance, 2, -1)
+    x1y1 = anchor_points - lt
+    x2y2 = anchor_points + rb
+    bbox = torch.cat([x1y1, x2y2], -1)
+    return bbox
+
+_m.dist2bbox.__code__ = _dist2bbox.__code__
+
+
+class DeepStreamOutput(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        boxes = x[:, :, :4]
+        objectness = x[:, :, 4:5]
+        scores, labels = torch.max(x[:, :, 5:], dim=-1, keepdim=True)
+        scores *= objectness
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
+
+
+def gold_yolo_export(weights, device, inplace=True, fuse=True):
+    model = load_checkpoint(weights, map_location=device, inplace=inplace, fuse=fuse)
+    model = switch_to_deploy(model)
+    for layer in model.modules():
+        t = type(layer)
+        if t.__name__ == 'RepVGGBlock':
+            layer.switch_to_deploy()
+    model.eval()
+    for k, m in model.named_modules():
+        if m.__class__.__name__ == 'Conv':
+            if isinstance(m.act, nn.SiLU):
+                m.act = SiLU()
+        elif m.__class__.__name__ == 'Detect':
+            m.inplace = False
+    return model
+
+
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
+    warnings.filterwarnings('ignore', category=UserWarning)
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
+
+
+def main(args):
+    suppress_warnings()
+
+    print(f'\nStarting: {args.weights}')
+
+    print('Opening Gold-YOLO model')
+
+    device = torch.device('cpu')
+    model = gold_yolo_export(args.weights, device)
+
+    model = nn.Sequential(model, DeepStreamOutput())
+
+    img_size = args.size * 2 if len(args.size) == 1 else args.size
+
+    onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
+    onnx_output_file = f'{args.weights}.onnx'
+
+    dynamic_axes = {
+        'input': {
+            0: 'batch'
+        },
+        'output': {
+            0: 'batch'
+        }
+    }
+
+    print('Exporting the model to ONNX')
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
+
+    if args.simplify:
+        print('Simplifying the ONNX model')
+        import onnxslim
+        model_onnx = onnx.load(onnx_output_file)
+        model_onnx = onnxslim.slim(model_onnx)
+        onnx.save(model_onnx, onnx_output_file)
+
+    print(f'Done: {onnx_output_file}\n')
+
+
+def parse_args():
+    import argparse
+    parser = argparse.ArgumentParser(description='DeepStream Gold-YOLO conversion')
+    parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
+    parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
+    parser.add_argument('--opset', type=int, default=13, help='ONNX opset version')
+    parser.add_argument('--simplify', action='store_true', help='ONNX simplify model')
+    parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size')
+    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('Invalid weights file')
+    if args.dynamic and args.batch > 1:
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
+    return args
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    main(args)

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)

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)

utils/export_rtdetr_pytorch.py

@@ -1,31 +1,28 @@
 import os
-import sys
-import argparse
-import warnings
 import onnx
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
+
 from src.core import YAMLConfig
 
 
 class DeepStreamOutput(nn.Module):
-    def __init__(self, img_size):
-        self.img_size = img_size
+    def __init__(self, img_size, use_focal_loss):
         super().__init__()
+        self.img_size = img_size
+        self.use_focal_loss = use_focal_loss
 
     def forward(self, x):
         boxes = x['pred_boxes']
-        boxes[:, :, [0, 2]] *= self.img_size[1]
-        boxes[:, :, [1, 3]] *= self.img_size[0]
-        scores, classes = torch.max(x['pred_logits'], 2, keepdim=True)
-        classes = classes.float()
-        return boxes, scores, classes
-
-
-def suppress_warnings():
-    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
-    warnings.filterwarnings('ignore', category=UserWarning)
-    warnings.filterwarnings('ignore', category=DeprecationWarning)
+        convert_matrix = torch.tensor(
+            [[1, 0, 1, 0], [0, 1, 0, 1], [-0.5, 0, 0.5, 0], [0, -0.5, 0, 0.5]], dtype=boxes.dtype, device=boxes.device
+        )
+        boxes @= convert_matrix
+        boxes *= torch.as_tensor([[*self.img_size]]).flip(1).tile([1, 2]).unsqueeze(1)
+        scores = F.sigmoid(x['pred_logits']) if self.use_focal_loss else F.softmax(x['pred_logits'])[:, :, :-1]
+        scores, labels = torch.max(scores, dim=-1, keepdim=True)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
 
 
 def rtdetr_pytorch_export(weights, cfg_file, device):
@@ -36,57 +33,62 @@ def rtdetr_pytorch_export(weights, cfg_file, device):
     else:
         state = checkpoint['model']
     cfg.model.load_state_dict(state)
-    return cfg.model.deploy()
+    return cfg.model.deploy(), cfg.postprocessor.use_focal_loss
+
+
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
+    warnings.filterwarnings('ignore', category=UserWarning)
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
 
 
 def main(args):
     suppress_warnings()
 
-    print('\nStarting: %s' % args.weights)
+    print(f'\nStarting: {args.weights}')
 
-    print('Opening RT-DETR PyTorch model\n')
+    print('Opening RT-DETR PyTorch model')
 
     device = torch.device('cpu')
-    model = rtdetr_pytorch_export(args.weights, args.config, device)
+    model, use_focal_loss = rtdetr_pytorch_export(args.weights, args.config, device)
 
     img_size = args.size * 2 if len(args.size) == 1 else args.size
 
-    model = nn.Sequential(model, DeepStreamOutput(img_size))
+    model = nn.Sequential(model, DeepStreamOutput(img_size, use_focal_loss))
 
     onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
-    onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx'
+    onnx_output_file = f'{args.weights}.onnx'
 
     dynamic_axes = {
         'input': {
             0: 'batch'
         },
-        'boxes': {
-            0: 'batch'
-        },
-        'scores': {
-            0: 'batch'
-        },
-        'classes': {
+        'output': {
             0: 'batch'
         }
     }
 
-    print('\nExporting the model to ONNX')
-    torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
-                      do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
-                      dynamic_axes=dynamic_axes if args.dynamic else None)
+    print('Exporting the model to ONNX')
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
 
     if args.simplify:
         print('Simplifying the ONNX model')
-        import onnxsim
+        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('Done: %s\n' % onnx_output_file)
+    print(f'Done: {onnx_output_file}\n')
 
 
 def parse_args():
+    import argparse
     parser = argparse.ArgumentParser(description='DeepStream RT-DETR PyTorch conversion')
     parser.add_argument('-w', '--weights', required=True, help='Input weights (.pth) file path (required)')
     parser.add_argument('-c', '--config', required=True, help='Input YAML (.yml) file path (required)')
@@ -107,4 +109,4 @@ def parse_args():
 
 if __name__ == '__main__':
     args = parse_args()
-    sys.exit(main(args))
+    main(args)

utils/export_rtdetr_ultralytics.py

@@ -1,34 +1,25 @@
 import os
-import sys
-import argparse
-import warnings
-import onnx
 import torch
 import torch.nn as nn
 from copy import deepcopy
+
 from ultralytics import RTDETR
-from ultralytics.utils.torch_utils import select_device
-from ultralytics.nn.modules import C2f, Detect, RTDETRDecoder
 
 
 class DeepStreamOutput(nn.Module):
     def __init__(self, img_size):
-        self.img_size = img_size
         super().__init__()
+        self.img_size = img_size
 
     def forward(self, x):
         boxes = x[:, :, :4]
-        boxes[:, :, [0, 2]] *= self.img_size[1]
-        boxes[:, :, [1, 3]] *= self.img_size[0]
-        scores, classes = torch.max(x[:, :, 4:], 2, keepdim=True)
-        classes = classes.float()
-        return boxes, scores, classes
-
-
-def suppress_warnings():
-    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
-    warnings.filterwarnings('ignore', category=UserWarning)
-    warnings.filterwarnings('ignore', category=DeprecationWarning)
+        convert_matrix = torch.tensor(
+            [[1, 0, 1, 0], [0, 1, 0, 1], [-0.5, 0, 0.5, 0], [0, -0.5, 0, 0.5]], dtype=boxes.dtype, device=boxes.device
+        )
+        boxes @= convert_matrix
+        boxes *= torch.as_tensor([[*self.img_size]]).flip(1).tile([1, 2]).unsqueeze(1)
+        scores, labels = torch.max(x[:, :, 4:], dim=-1, keepdim=True)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
 
 
 def rtdetr_ultralytics_export(weights, device):
@@ -40,74 +31,74 @@ def rtdetr_ultralytics_export(weights, device):
     model.float()
     model = model.fuse()
     for k, m in model.named_modules():
-        if isinstance(m, (Detect, RTDETRDecoder)):
+        if m.__class__.__name__ in ('Detect', 'RTDETRDecoder'):
             m.dynamic = False
             m.export = True
             m.format = 'onnx'
-        elif isinstance(m, C2f):
+        elif m.__class__.__name__ == 'C2f':
             m.forward = m.forward_split
     return model
 
 
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
+    warnings.filterwarnings('ignore', category=UserWarning)
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
+
+
 def main(args):
     suppress_warnings()
 
-    print('\nStarting: %s' % args.weights)
+    print(f'\nStarting: {args.weights}')
 
-    print('Opening RT-DETR Ultralytics model\n')
+    print('Opening RT-DETR Ultralytics model')
 
-    device = select_device('cpu')
+    device = torch.device('cpu')
     model = rtdetr_ultralytics_export(args.weights, device)
 
     if len(model.names.keys()) > 0:
-        print('\nCreating labels.txt file')
-        f = open('labels.txt', 'w')
-        for name in model.names.values():
-            f.write(name + '\n')
-        f.close()
+        print('Creating labels.txt file')
+        with open('labels.txt', 'w', encoding='utf-8') as f:
+            for name in model.names.values():
+                f.write(f'{name}\n')
 
     img_size = args.size * 2 if len(args.size) == 1 else args.size
 
     model = nn.Sequential(model, DeepStreamOutput(img_size))
 
     onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
-    onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx'
+    onnx_output_file = f'{args.weights}.onnx'
 
     dynamic_axes = {
         'input': {
             0: 'batch'
         },
-        'boxes': {
-            0: 'batch'
-        },
-        'scores': {
-            0: 'batch'
-        },
-        'classes': {
+        'output': {
             0: 'batch'
         }
     }
 
-    print('\nExporting the model to ONNX')
-    torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
-                      do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
-                      dynamic_axes=dynamic_axes if args.dynamic else None)
+    print('Exporting the model to ONNX')
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
 
     if args.simplify:
-        print('Simplifying the ONNX model')
-        import onnxsim
-        model_onnx = onnx.load(onnx_output_file)
-        model_onnx, _ = onnxsim.simplify(model_onnx)
-        onnx.save(model_onnx, onnx_output_file)
+        print('Simplifying is not available for this model')
 
-    print('Done: %s\n' % onnx_output_file)
+    print(f'Done: {onnx_output_file}\n')
 
 
 def parse_args():
+    import argparse
     parser = argparse.ArgumentParser(description='DeepStream RT-DETR Ultralytics conversion')
     parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
     parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
-    parser.add_argument('--opset', type=int, default=16, help='ONNX opset version')
+    parser.add_argument('--opset', type=int, default=17, help='ONNX opset version')
     parser.add_argument('--simplify', action='store_true', help='ONNX simplify model')
     parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size')
     parser.add_argument('--batch', type=int, default=1, help='Static batch-size')
@@ -121,4 +112,4 @@ def parse_args():
 
 if __name__ == '__main__':
     args = parse_args()
-    sys.exit(main(args))
+    main(args)

utils/export_rtmdet.py

@@ -0,0 +1,150 @@
+import os
+import types
+import onnx
+import torch
+import torch.nn as nn
+
+from mmdet.apis import init_detector
+from projects.easydeploy.model import DeployModel, MMYOLOBackend
+from projects.easydeploy.bbox_code import rtmdet_bbox_decoder as bbox_decoder
+
+
+class DeepStreamOutput(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        boxes = x[0]
+        scores, labels = torch.max(x[1], dim=-1, keepdim=True)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
+
+
+def pred_by_feat_deepstream(self, cls_scores, bbox_preds, objectnesses=None, **kwargs):
+    assert len(cls_scores) == len(bbox_preds)
+    dtype = cls_scores[0].dtype
+    device = cls_scores[0].device
+
+    num_imgs = cls_scores[0].shape[0]
+    featmap_sizes = [cls_score.shape[2:] for cls_score in cls_scores]
+
+    mlvl_priors = self.prior_generate(featmap_sizes, dtype=dtype, device=device)
+
+    flatten_priors = torch.cat(mlvl_priors)
+
+    mlvl_strides = [
+        flatten_priors.new_full(
+            (featmap_size[0] * featmap_size[1] * self.num_base_priors,), stride
+        ) for featmap_size, stride in zip(
+            featmap_sizes, self.featmap_strides
+        )
+    ]
+    flatten_stride = torch.cat(mlvl_strides)
+
+    flatten_cls_scores = [
+        cls_score.permute(0, 2, 3, 1).reshape(num_imgs, -1, self.num_classes) for cls_score in cls_scores
+    ]
+    cls_scores = torch.cat(flatten_cls_scores, dim=1).sigmoid()
+
+    flatten_bbox_preds = [bbox_pred.permute(0, 2, 3, 1).reshape(num_imgs, -1, 4) for bbox_pred in bbox_preds]
+    flatten_bbox_preds = torch.cat(flatten_bbox_preds, dim=1)
+
+    if objectnesses is not None:
+        flatten_objectness = [objectness.permute(0, 2, 3, 1).reshape(num_imgs, -1) for objectness in objectnesses]
+        flatten_objectness = torch.cat(flatten_objectness, dim=1).sigmoid()
+        cls_scores = cls_scores * (flatten_objectness.unsqueeze(-1))
+
+    scores = cls_scores
+
+    bboxes = bbox_decoder(flatten_priors[None], flatten_bbox_preds, flatten_stride)
+
+    return bboxes, scores
+
+
+def rtmdet_export(weights, config, device):
+    model = init_detector(config, weights, device=device)
+    model.eval()
+    deploy_model = DeployModel(baseModel=model, backend=MMYOLOBackend.ONNXRUNTIME, postprocess_cfg=None)
+    deploy_model.eval()
+    deploy_model.with_postprocess = True
+    deploy_model.prior_generate = model.bbox_head.prior_generator.grid_priors
+    deploy_model.num_base_priors = model.bbox_head.num_base_priors
+    deploy_model.featmap_strides = model.bbox_head.featmap_strides
+    deploy_model.num_classes = model.bbox_head.num_classes
+    deploy_model.pred_by_feat = types.MethodType(pred_by_feat_deepstream, deploy_model)
+    return deploy_model
+
+
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
+    warnings.filterwarnings('ignore', category=UserWarning)
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
+
+
+def main(args):
+    suppress_warnings()
+
+    print(f'\nStarting: {args.weights}')
+
+    print('Opening RTMDet model')
+
+    device = torch.device('cpu')
+    model = rtmdet_export(args.weights, args.config, device)
+
+    model = nn.Sequential(model, DeepStreamOutput())
+
+    img_size = args.size * 2 if len(args.size) == 1 else args.size
+
+    onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
+    onnx_output_file = f'{args.weights}.onnx'
+
+    dynamic_axes = {
+        'input': {
+            0: 'batch'
+        },
+        'output': {
+            0: 'batch'
+        }
+    }
+
+    print('Exporting the model to ONNX')
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
+
+    if args.simplify:
+        print('Simplifying the ONNX model')
+        import onnxslim
+        model_onnx = onnx.load(onnx_output_file)
+        model_onnx = onnxslim.slim(model_onnx)
+        onnx.save(model_onnx, onnx_output_file)
+
+    print(f'Done: {onnx_output_file}\n')
+
+
+def parse_args():
+    import argparse
+    parser = argparse.ArgumentParser(description='DeepStream RTMDet conversion')
+    parser.add_argument('-w', '--weights', required=True, type=str, help='Input weights (.pt) file path (required)')
+    parser.add_argument('-c', '--config', required=True, help='Input config (.py) file path (required)')
+    parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
+    parser.add_argument('--opset', type=int, default=17, help='ONNX opset version')
+    parser.add_argument('--simplify', action='store_true', help='ONNX simplify model')
+    parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size')
+    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('Invalid weights file')
+    if not os.path.isfile(args.config):
+        raise SystemExit('Invalid config file')
+    if args.dynamic and args.batch > 1:
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
+    return args
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    main(args)

utils/export_yoloV5.py

@@ -1,13 +1,9 @@
 import os
-import sys
-import argparse
-import warnings
 import onnx
 import torch
 import torch.nn as nn
+
 from models.experimental import attempt_load
-from utils.torch_utils import select_device
-from models.yolo import Detect
 
 
 class DeepStreamOutput(nn.Module):
@@ -17,46 +13,51 @@ class DeepStreamOutput(nn.Module):
     def forward(self, x):
         x = x[0]
         boxes = x[:, :, :4]
+        convert_matrix = torch.tensor(
+            [[1, 0, 1, 0], [0, 1, 0, 1], [-0.5, 0, 0.5, 0], [0, -0.5, 0, 0.5]], dtype=boxes.dtype, device=boxes.device
+        )
+        boxes @= convert_matrix
         objectness = x[:, :, 4:5]
-        scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
+        scores, labels = torch.max(x[:, :, 5:], dim=-1, keepdim=True)
         scores *= objectness
-        classes = classes.float()
-        return boxes, scores, classes
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
 
 
-def suppress_warnings():
-    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
-    warnings.filterwarnings('ignore', category=UserWarning)
-    warnings.filterwarnings('ignore', category=DeprecationWarning)
-
-
-def yolov5_export(weights, device):
-    model = attempt_load(weights, device=device, inplace=True, fuse=True)
+def yolov5_export(weights, device, inplace=True, fuse=True):
+    model = attempt_load(weights, device=device, inplace=inplace, fuse=fuse)
     model.eval()
     for k, m in model.named_modules():
-        if isinstance(m, Detect):
+        if m.__class__.__name__ == 'Detect':
             m.inplace = False
             m.dynamic = False
             m.export = True
     return model
 
 
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
+    warnings.filterwarnings('ignore', category=UserWarning)
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
+
+
 def main(args):
     suppress_warnings()
 
-    print('\nStarting: %s' % args.weights)
+    print(f'\nStarting: {args.weights}')
 
-    print('Opening YOLOv5 model\n')
+    print('Opening YOLOv5 model')
 
-    device = select_device('cpu')
+    device = torch.device('cpu')
     model = yolov5_export(args.weights, device)
 
     if len(model.names.keys()) > 0:
-        print('\nCreating labels.txt file')
-        f = open('labels.txt', 'w')
-        for name in model.names.values():
-            f.write(name + '\n')
-        f.close()
+        print('Creating labels.txt file')
+        with open('labels.txt', 'w', encoding='utf-8') as f:
+            for name in model.names.values():
+                f.write(f'{name}\n')
 
     model = nn.Sequential(model, DeepStreamOutput())
 
@@ -66,41 +67,37 @@ def main(args):
         img_size = [1280] * 2
 
     onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
-    onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx'
+    onnx_output_file = f'{args.weights}.onnx'
 
     dynamic_axes = {
         'input': {
             0: 'batch'
         },
-        'boxes': {
-            0: 'batch'
-        },
-        'scores': {
-            0: 'batch'
-        },
-        'classes': {
+        'output': {
             0: 'batch'
         }
     }
 
-    print('\nExporting the model to ONNX')
-    torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
-                      do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
-                      dynamic_axes=dynamic_axes if args.dynamic else None)
+    print('Exporting the model to ONNX')
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
 
     if args.simplify:
         print('Simplifying the ONNX model')
-        import onnxsim
+        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('Done: %s\n' % onnx_output_file)
+    print(f'Done: {onnx_output_file}\n')
 
 
 def parse_args():
+    import argparse
     parser = argparse.ArgumentParser(description='DeepStream YOLOv5 conversion')
-    parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
+    parser.add_argument('-w', '--weights', required=True, type=str, help='Input weights (.pt) file path (required)')
     parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
     parser.add_argument('--p6', action='store_true', help='P6 model')
     parser.add_argument('--opset', type=int, default=17, help='ONNX opset version')
@@ -117,4 +114,4 @@ def parse_args():
 
 if __name__ == '__main__':
     args = parse_args()
-    sys.exit(main(args))
+    main(args)

utils/export_yoloV6.py

@@ -1,13 +1,11 @@
 import os
-import sys
-import argparse
-import warnings
 import onnx
 import torch
 import torch.nn as nn
-from yolov6.utils.checkpoint import load_checkpoint
-from yolov6.layers.common import RepVGGBlock, SiLU
+
 from yolov6.models.effidehead import Detect
+from yolov6.layers.common import RepVGGBlock, SiLU
+from yolov6.utils.checkpoint import load_checkpoint
 
 try:
     from yolov6.layers.common import ConvModule
@@ -21,17 +19,14 @@ class DeepStreamOutput(nn.Module):
 
     def forward(self, x):
         boxes = x[:, :, :4]
+        convert_matrix = torch.tensor(
+            [[1, 0, 1, 0], [0, 1, 0, 1], [-0.5, 0, 0.5, 0], [0, -0.5, 0, 0.5]], dtype=boxes.dtype, device=boxes.device
+        )
+        boxes @= convert_matrix
         objectness = x[:, :, 4:5]
-        scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
+        scores, labels = torch.max(x[:, :, 5:], dim=-1, keepdim=True)
         scores *= objectness
-        classes = classes.float()
-        return boxes, scores, classes
-
-
-def suppress_warnings():
-    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
-    warnings.filterwarnings('ignore', category=UserWarning)
-    warnings.filterwarnings('ignore', category=DeprecationWarning)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
 
 
 def yolov6_export(weights, device):
@@ -51,12 +46,21 @@ def yolov6_export(weights, device):
     return model
 
 
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
+    warnings.filterwarnings('ignore', category=UserWarning)
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
+
+
 def main(args):
     suppress_warnings()
 
-    print('\nStarting: %s' % args.weights)
+    print(f'\nStarting: {args.weights}')
 
-    print('Opening YOLOv6 model\n')
+    print('Opening YOLOv6 model')
 
     device = torch.device('cpu')
     model = yolov6_export(args.weights, device)
@@ -69,39 +73,35 @@ def main(args):
         img_size = [1280] * 2
 
     onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
-    onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx'
+    onnx_output_file = f'{args.weights}.onnx'
 
     dynamic_axes = {
         'input': {
             0: 'batch'
         },
-        'boxes': {
-            0: 'batch'
-        },
-        'scores': {
-            0: 'batch'
-        },
-        'classes': {
+        'output': {
             0: 'batch'
         }
     }
 
-    print('\nExporting the model to ONNX')
-    torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
-                      do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
-                      dynamic_axes=dynamic_axes if args.dynamic else None)
+    print('Exporting the model to ONNX')
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
 
     if args.simplify:
         print('Simplifying the ONNX model')
-        import onnxsim
+        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('Done: %s\n' % onnx_output_file)
+    print(f'Done: {onnx_output_file}\n')
 
 
 def parse_args():
+    import argparse
     parser = argparse.ArgumentParser(description='DeepStream YOLOv6 conversion')
     parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
     parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
@@ -120,4 +120,4 @@ def parse_args():
 
 if __name__ == '__main__':
     args = parse_args()
-    sys.exit(main(args))
+    main(args)

utils/export_yoloV7.py

@@ -1,10 +1,8 @@
 import os
-import sys
-import argparse
-import warnings
 import onnx
 import torch
 import torch.nn as nn
+
 import models
 from models.experimental import attempt_load
 from utils.torch_utils import select_device
@@ -17,17 +15,14 @@ class DeepStreamOutput(nn.Module):
 
     def forward(self, x):
         boxes = x[:, :, :4]
+        convert_matrix = torch.tensor(
+            [[1, 0, 1, 0], [0, 1, 0, 1], [-0.5, 0, 0.5, 0], [0, -0.5, 0, 0.5]], dtype=boxes.dtype, device=boxes.device
+        )
+        boxes @= convert_matrix
         objectness = x[:, :, 4:5]
-        scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
+        scores, labels = torch.max(x[:, :, 5:], dim=-1, keepdim=True)
         scores *= objectness
-        classes = classes.float()
-        return boxes, scores, classes
-
-
-def suppress_warnings():
-    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
-    warnings.filterwarnings('ignore', category=UserWarning)
-    warnings.filterwarnings('ignore', category=DeprecationWarning)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
 
 
 def yolov7_export(weights, device):
@@ -45,22 +40,30 @@ def yolov7_export(weights, device):
     return model
 
 
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
+    warnings.filterwarnings('ignore', category=UserWarning)
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
+
+
 def main(args):
     suppress_warnings()
 
-    print('\nStarting: %s' % args.weights)
+    print(f'\nStarting: {args.weights}')
 
-    print('Opening YOLOv7 model\n')
+    print('Opening YOLOv7 model')
 
     device = select_device('cpu')
     model = yolov7_export(args.weights, device)
 
-    if len(model.names) > 0:
-        print('\nCreating labels.txt file')
-        f = open('labels.txt', 'w')
-        for name in model.names:
-            f.write(name + '\n')
-        f.close()
+    if hasattr(model, 'names') and len(model.names) > 0:
+        print('Creating labels.txt file')
+        with open('labels.txt', 'w', encoding='utf-8') as f:
+            for name in model.names:
+                f.write(f'{name}\n')
 
     model = nn.Sequential(model, DeepStreamOutput())
 
@@ -70,39 +73,35 @@ def main(args):
         img_size = [1280] * 2
 
     onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
-    onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx'
+    onnx_output_file = f'{args.weights}.onnx'
 
     dynamic_axes = {
         'input': {
             0: 'batch'
         },
-        'boxes': {
-            0: 'batch'
-        },
-        'scores': {
-            0: 'batch'
-        },
-        'classes': {
+        'output': {
             0: 'batch'
         }
     }
 
-    print('\nExporting the model to ONNX')
-    torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
-                      do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
-                      dynamic_axes=dynamic_axes if args.dynamic else None)
+    print('Exporting the model to ONNX')
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
 
     if args.simplify:
         print('Simplifying the ONNX model')
-        import onnxsim
+        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('Done: %s\n' % onnx_output_file)
+    print(f'Done: {onnx_output_file}\n')
 
 
 def parse_args():
+    import argparse
     parser = argparse.ArgumentParser(description='DeepStream YOLOv7 conversion')
     parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
     parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
@@ -121,4 +120,4 @@ def parse_args():
 
 if __name__ == '__main__':
     args = parse_args()
-    sys.exit(main(args))
+    main(args)

utils/export_yoloV7_u6.py

@@ -1,13 +1,11 @@
 import os
-import sys
-import argparse
-import warnings
 import onnx
 import torch
 import torch.nn as nn
+
+from utils.torch_utils import select_device
 from models.experimental import attempt_load
 from models.yolo import Detect, V6Detect, IV6Detect
-from utils.torch_utils import select_device
 
 
 class DeepStreamOutput(nn.Module):
@@ -17,15 +15,12 @@ class DeepStreamOutput(nn.Module):
     def forward(self, x):
         x = x.transpose(1, 2)
         boxes = x[:, :, :4]
-        scores, classes = torch.max(x[:, :, 4:], 2, keepdim=True)
-        classes = classes.float()
-        return boxes, scores, classes
-
-
-def suppress_warnings():
-    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
-    warnings.filterwarnings('ignore', category=UserWarning)
-    warnings.filterwarnings('ignore', category=DeprecationWarning)
+        convert_matrix = torch.tensor(
+            [[1, 0, 1, 0], [0, 1, 0, 1], [-0.5, 0, 0.5, 0], [0, -0.5, 0, 0.5]], dtype=boxes.dtype, device=boxes.device
+        )
+        boxes @= convert_matrix
+        scores, labels = torch.max(x[:, :, 4:], dim=-1, keepdim=True)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
 
 
 def yolov7_u6_export(weights, device):
@@ -39,61 +34,65 @@ def yolov7_u6_export(weights, device):
     return model
 
 
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
+    warnings.filterwarnings('ignore', category=UserWarning)
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
+
+
 def main(args):
     suppress_warnings()
 
-    print('\nStarting: %s' % args.weights)
+    print(f'\nStarting: {args.weights}')
 
-    print('Opening YOLOv7_u6 model\n')
+    print('Opening YOLOv7_u6 model')
 
     device = select_device('cpu')
     model = yolov7_u6_export(args.weights, device)
 
     if len(model.names.keys()) > 0:
-        print('\nCreating labels.txt file')
-        f = open('labels.txt', 'w')
-        for name in model.names.values():
-            f.write(name + '\n')
-        f.close()
+        print('Creating labels.txt file')
+        with open('labels.txt', 'w', encoding='utf-8') as f:
+            for name in model.names.values():
+                f.write(f'{name}\n')
 
     model = nn.Sequential(model, DeepStreamOutput())
 
     img_size = args.size * 2 if len(args.size) == 1 else args.size
 
     onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
-    onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx'
+    onnx_output_file = f'{args.weights}.onnx'
 
     dynamic_axes = {
         'input': {
             0: 'batch'
         },
-        'boxes': {
-            0: 'batch'
-        },
-        'scores': {
-            0: 'batch'
-        },
-        'classes': {
+        'output': {
             0: 'batch'
         }
     }
 
-    print('\nExporting the model to ONNX')
-    torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
-                      do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
-                      dynamic_axes=dynamic_axes if args.dynamic else None)
+    print('Exporting the model to ONNX')
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
 
     if args.simplify:
         print('Simplifying the ONNX model')
-        import onnxsim
+        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('Done: %s\n' % onnx_output_file)
+    print(f'Done: {onnx_output_file}\n')
 
 
 def parse_args():
+    import argparse
     parser = argparse.ArgumentParser(description='DeepStream YOLOv7-u6 conversion')
     parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
     parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
@@ -111,4 +110,4 @@ def parse_args():
 
 if __name__ == '__main__':
     args = parse_args()
-    sys.exit(main(args))
+    main(args)

utils/export_yoloV8.py

@@ -1,14 +1,26 @@
 import os
 import sys
-import argparse
-import warnings
 import onnx
 import torch
 import torch.nn as nn
 from copy import deepcopy
-from ultralytics import YOLO
-from ultralytics.utils.torch_utils import select_device
-from ultralytics.nn.modules import C2f, Detect, RTDETRDecoder
+
+import ultralytics.utils
+import ultralytics.models.yolo
+import ultralytics.utils.tal as _m
+
+sys.modules['ultralytics.yolo'] = ultralytics.models.yolo
+sys.modules['ultralytics.yolo.utils'] = ultralytics.utils
+
+
+def _dist2bbox(distance, anchor_points, xywh=False, dim=-1):
+    lt, rb = distance.chunk(2, dim)
+    x1y1 = anchor_points - lt
+    x2y2 = anchor_points + rb
+    return torch.cat((x1y1, x2y2), dim)
+
+
+_m.dist2bbox.__code__ = _dist2bbox.__code__
 
 
 class DeepStreamOutput(nn.Module):
@@ -18,94 +30,103 @@ class DeepStreamOutput(nn.Module):
     def forward(self, x):
         x = x.transpose(1, 2)
         boxes = x[:, :, :4]
-        scores, classes = torch.max(x[:, :, 4:], 2, keepdim=True)
-        classes = classes.float()
-        return boxes, scores, classes
+        scores, labels = torch.max(x[:, :, 4:], dim=-1, keepdim=True)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
 
 
-def suppress_warnings():
-    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
-    warnings.filterwarnings('ignore', category=UserWarning)
-    warnings.filterwarnings('ignore', category=DeprecationWarning)
-
-
-def yolov8_export(weights, device):
-    model = YOLO(weights)
-    model = deepcopy(model.model).to(device)
+def yolov8_export(weights, device, inplace=True, fuse=True):
+    ckpt = torch.load(weights, map_location='cpu')
+    ckpt = (ckpt.get('ema') or ckpt['model']).to(device).float()
+    if not hasattr(ckpt, 'stride'):
+        ckpt.stride = torch.tensor([32.])
+    if hasattr(ckpt, 'names') and isinstance(ckpt.names, (list, tuple)):
+        ckpt.names = dict(enumerate(ckpt.names))
+    model = ckpt.fuse().eval() if fuse and hasattr(ckpt, 'fuse') else ckpt.eval()
+    for m in model.modules():
+        t = type(m)
+        if hasattr(m, 'inplace'):
+            m.inplace = inplace
+        elif t.__name__ == 'Upsample' and not hasattr(m, 'recompute_scale_factor'):
+            m.recompute_scale_factor = None
+    model = deepcopy(model).to(device)
     for p in model.parameters():
         p.requires_grad = False
     model.eval()
     model.float()
     model = model.fuse()
     for k, m in model.named_modules():
-        if isinstance(m, (Detect, RTDETRDecoder)):
+        if m.__class__.__name__ in ('Detect', 'RTDETRDecoder'):
             m.dynamic = False
             m.export = True
             m.format = 'onnx'
-        elif isinstance(m, C2f):
+        elif m.__class__.__name__ == 'C2f':
             m.forward = m.forward_split
     return model
 
 
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
+    warnings.filterwarnings('ignore', category=UserWarning)
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
+
+
 def main(args):
     suppress_warnings()
 
-    print('\nStarting: %s' % args.weights)
+    print(f'\nStarting: {args.weights}')
 
-    print('Opening YOLOv8 model\n')
+    print('Opening YOLOv8 model')
 
-    device = select_device('cpu')
+    device = torch.device('cpu')
     model = yolov8_export(args.weights, device)
 
     if len(model.names.keys()) > 0:
-        print('\nCreating labels.txt file')
-        f = open('labels.txt', 'w')
-        for name in model.names.values():
-            f.write(name + '\n')
-        f.close()
+        print('Creating labels.txt file')
+        with open('labels.txt', 'w', encoding='utf-8') as f:
+            for name in model.names.values():
+                f.write(f'{name}\n')
 
     model = nn.Sequential(model, DeepStreamOutput())
 
     img_size = args.size * 2 if len(args.size) == 1 else args.size
 
     onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
-    onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx'
+    onnx_output_file = f'{args.weights}.onnx'
 
     dynamic_axes = {
         'input': {
             0: 'batch'
         },
-        'boxes': {
-            0: 'batch'
-        },
-        'scores': {
-            0: 'batch'
-        },
-        'classes': {
+        'output': {
             0: 'batch'
         }
     }
 
-    print('\nExporting the model to ONNX')
-    torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
-                      do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
-                      dynamic_axes=dynamic_axes if args.dynamic else None)
+    print('Exporting the model to ONNX')
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
 
     if args.simplify:
         print('Simplifying the ONNX model')
-        import onnxsim
+        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('Done: %s\n' % onnx_output_file)
+    print(f'Done: {onnx_output_file}\n')
 
 
 def parse_args():
+    import argparse
     parser = argparse.ArgumentParser(description='DeepStream YOLOv8 conversion')
     parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
     parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
-    parser.add_argument('--opset', type=int, default=16, help='ONNX opset version')
+    parser.add_argument('--opset', type=int, default=17, help='ONNX opset version')
     parser.add_argument('--simplify', action='store_true', help='ONNX simplify model')
     parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size')
     parser.add_argument('--batch', type=int, default=1, help='Static batch-size')
@@ -119,4 +140,4 @@ def parse_args():
 
 if __name__ == '__main__':
     args = parse_args()
-    sys.exit(main(args))
+    main(args)

utils/export_yoloV9.py

@@ -0,0 +1,145 @@
+import os
+import onnx
+import torch
+import torch.nn as nn
+
+import utils.tal.anchor_generator as _m
+
+
+def _dist2bbox(distance, anchor_points, xywh=False, dim=-1):
+    lt, rb = torch.split(distance, 2, dim)
+    x1y1 = anchor_points - lt
+    x2y2 = anchor_points + rb
+    return torch.cat((x1y1, x2y2), dim)
+
+
+_m.dist2bbox.__code__ = _dist2bbox.__code__
+
+
+class DeepStreamOutputDual(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        x = x[1].transpose(1, 2)
+        boxes = x[:, :, :4]
+        scores, labels = torch.max(x[:, :, 4:], dim=-1, keepdim=True)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
+
+
+class DeepStreamOutput(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        x = x.transpose(1, 2)
+        boxes = x[:, :, :4]
+        scores, labels = torch.max(x[:, :, 4:], dim=-1, keepdim=True)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
+
+
+def yolov9_export(weights, device, inplace=True, fuse=True):
+    ckpt = torch.load(weights, map_location='cpu')
+    ckpt = (ckpt.get('ema') or ckpt['model']).to(device).float()
+    if not hasattr(ckpt, 'stride'):
+        ckpt.stride = torch.tensor([32.])
+    if hasattr(ckpt, 'names') and isinstance(ckpt.names, (list, tuple)):
+        ckpt.names = dict(enumerate(ckpt.names))
+    model = ckpt.fuse().eval() if fuse and hasattr(ckpt, 'fuse') else ckpt.eval()
+    for m in model.modules():
+        t = type(m)
+        if t.__name__ in ('Hardswish', 'LeakyReLU', 'ReLU', 'ReLU6', 'SiLU', 'Detect', 'Model'):
+            m.inplace = inplace
+        elif t.__name__ == 'Upsample' and not hasattr(m, 'recompute_scale_factor'):
+            m.recompute_scale_factor = None
+    model.eval()
+    head = 'Detect'
+    for k, m in model.named_modules():
+        if m.__class__.__name__ in ('Detect', 'DDetect', 'DualDetect', 'DualDDetect'):
+            m.inplace = False
+            m.dynamic = False
+            m.export = True
+            head = m.__class__.__name__
+    return model, head
+
+
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
+    warnings.filterwarnings('ignore', category=UserWarning)
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
+
+
+def main(args):
+    suppress_warnings()
+
+    print(f'\nStarting: {args.weights}')
+
+    print('Opening YOLOv9 model')
+
+    device = torch.device('cpu')
+    model, head = yolov9_export(args.weights, device)
+
+    if len(model.names.keys()) > 0:
+        print('Creating labels.txt file')
+        with open('labels.txt', 'w', encoding='utf-8') as f:
+            for name in model.names.values():
+                f.write(f'{name}\n')
+
+    if head in ('Detect', 'DDetect'):
+        model = nn.Sequential(model, DeepStreamOutput())
+    else:
+        model = nn.Sequential(model, DeepStreamOutputDual())
+
+    img_size = args.size * 2 if len(args.size) == 1 else args.size
+
+    onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
+    onnx_output_file = f'{args.weights}.onnx'
+
+    dynamic_axes = {
+        'input': {
+            0: 'batch'
+        },
+        'output': {
+            0: 'batch'
+        }
+    }
+
+    print('Exporting the model to ONNX')
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
+
+    if args.simplify:
+        print('Simplifying the ONNX model')
+        import onnxslim
+        model_onnx = onnx.load(onnx_output_file)
+        model_onnx = onnxslim.slim(model_onnx)
+        onnx.save(model_onnx, onnx_output_file)
+
+    print(f'Done: {onnx_output_file}\n')
+
+
+def parse_args():
+    import argparse
+    parser = argparse.ArgumentParser(description='DeepStream YOLOv9 conversion')
+    parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
+    parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
+    parser.add_argument('--opset', type=int, default=17, help='ONNX opset version')
+    parser.add_argument('--simplify', action='store_true', help='ONNX simplify model')
+    parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size')
+    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('Invalid weights file')
+    if args.dynamic and args.batch > 1:
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
+    return args
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    main(args)

utils/export_yolonas.py

@@ -1,10 +1,8 @@
 import os
-import sys
-import argparse
-import warnings
 import onnx
 import torch
 import torch.nn as nn
+
 from super_gradients.training import models
 
 
@@ -14,15 +12,17 @@ class DeepStreamOutput(nn.Module):
 
     def forward(self, x):
         boxes = x[0]
-        scores, classes = torch.max(x[1], 2, keepdim=True)
-        classes = classes.float()
-        return boxes, scores, classes
+        scores, labels = torch.max(x[1], dim=-1, keepdim=True)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
 
 
 def suppress_warnings():
+    import warnings
     warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
     warnings.filterwarnings('ignore', category=UserWarning)
     warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
 
 
 def yolonas_export(model_name, weights, num_classes, size):
@@ -36,9 +36,9 @@ def yolonas_export(model_name, weights, num_classes, size):
 def main(args):
     suppress_warnings()
 
-    print('\nStarting: %s' % args.weights)
+    print(f'\nStarting: {args.weights}')
 
-    print('Opening YOLO-NAS model\n')
+    print('Opening YOLO-NAS model')
 
     device = torch.device('cpu')
     model = yolonas_export(args.model, args.weights, args.classes, args.size)
@@ -48,39 +48,35 @@ def main(args):
     img_size = args.size * 2 if len(args.size) == 1 else args.size
 
     onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
-    onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx'
+    onnx_output_file = f'{args.weights}.onnx'
 
     dynamic_axes = {
         'input': {
             0: 'batch'
         },
-        'boxes': {
-            0: 'batch'
-        },
-        'scores': {
-            0: 'batch'
-        },
-        'classes': {
+        'output': {
             0: 'batch'
         }
     }
 
-    print('\nExporting the model to ONNX')
-    torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
-                      do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
-                      dynamic_axes=dynamic_axes if args.dynamic else None)
+    print('Exporting the model to ONNX')
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
 
     if args.simplify:
         print('Simplifying the ONNX model')
-        import onnxsim
+        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('Done: %s\n' % onnx_output_file)
+    print(f'Done: {onnx_output_file}\n')
 
 
 def parse_args():
+    import argparse
     parser = argparse.ArgumentParser(description='DeepStream YOLO-NAS conversion')
     parser.add_argument('-m', '--model', required=True, help='Model name (required)')
     parser.add_argument('-w', '--weights', required=True, help='Input weights (.pth) file path (required)')
@@ -102,4 +98,4 @@ def parse_args():
 
 if __name__ == '__main__':
     args = parse_args()
-    sys.exit(main(args))
+    main(args)

utils/export_yolor.py

@@ -1,7 +1,4 @@
 import os
-import sys
-import argparse
-import warnings
 import onnx
 import torch
 import torch.nn as nn
@@ -14,17 +11,14 @@ class DeepStreamOutput(nn.Module):
     def forward(self, x):
         x = x[0]
         boxes = x[:, :, :4]
+        convert_matrix = torch.tensor(
+            [[1, 0, 1, 0], [0, 1, 0, 1], [-0.5, 0, 0.5, 0], [0, -0.5, 0, 0.5]], dtype=boxes.dtype, device=boxes.device
+        )
+        boxes @= convert_matrix
         objectness = x[:, :, 4:5]
-        scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
+        scores, labels = torch.max(x[:, :, 5:], dim=-1, keepdim=True)
         scores *= objectness
-        classes = classes.float()
-        return boxes, scores, classes
-
-
-def suppress_warnings():
-    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
-    warnings.filterwarnings('ignore', category=UserWarning)
-    warnings.filterwarnings('ignore', category=DeprecationWarning)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
 
 
 def yolor_export(weights, cfg, size, device):
@@ -57,22 +51,30 @@ def yolor_export(weights, cfg, size, device):
     return model
 
 
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
+    warnings.filterwarnings('ignore', category=UserWarning)
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
+
+
 def main(args):
     suppress_warnings()
 
-    print('\nStarting: %s' % args.weights)
+    print(f'\nStarting: {args.weights}')
 
-    print('Opening YOLOR model\n')
+    print('Opening YOLOR model')
 
     device = torch.device('cpu')
     model = yolor_export(args.weights, args.cfg, args.size, device)
 
     if hasattr(model, 'names') and len(model.names) > 0:
-        print('\nCreating labels.txt file')
-        f = open('labels.txt', 'w')
-        for name in model.names:
-            f.write(name + '\n')
-        f.close()
+        print('Creating labels.txt file')
+        with open('labels.txt', 'w', encoding='utf-8') as f:
+            for name in model.names:
+                f.write(f'{name}\n')
 
     model = nn.Sequential(model, DeepStreamOutput())
 
@@ -82,41 +84,37 @@ def main(args):
         img_size = [1280] * 2
 
     onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
-    onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx'
+    onnx_output_file = f'{args.weights}.onnx'
 
     dynamic_axes = {
         'input': {
             0: 'batch'
         },
-        'boxes': {
-            0: 'batch'
-        },
-        'scores': {
-            0: 'batch'
-        },
-        'classes': {
+        'output': {
             0: 'batch'
         }
     }
 
-    print('\nExporting the model to ONNX')
-    torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
-                      do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
-                      dynamic_axes=dynamic_axes if args.dynamic else None)
+    print('Exporting the model to ONNX')
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
 
     if args.simplify:
         print('Simplifying the ONNX model')
-        import onnxsim
+        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('Done: %s\n' % onnx_output_file)
+    print(f'Done: {onnx_output_file}\n')
 
 
 def parse_args():
+    import argparse
     parser = argparse.ArgumentParser(description='DeepStream YOLOR conversion')
-    parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
+    parser.add_argument('-w', '--weights', required=True, type=str, help='Input weights (.pt) file path (required)')
     parser.add_argument('-c', '--cfg', default='', help='Input cfg (.cfg) file path')
     parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
     parser.add_argument('--p6', action='store_true', help='P6 model')
@@ -134,4 +132,4 @@ def parse_args():
 
 if __name__ == '__main__':
     args = parse_args()
-    sys.exit(main(args))
+    main(args)

utils/export_yolox.py

@@ -1,10 +1,8 @@
 import os
-import sys
-import argparse
-import warnings
 import onnx
 import torch
 import torch.nn as nn
+
 from yolox.exp import get_exp
 from yolox.utils import replace_module
 from yolox.models.network_blocks import SiLU
@@ -16,17 +14,14 @@ class DeepStreamOutput(nn.Module):
 
     def forward(self, x):
         boxes = x[:, :, :4]
+        convert_matrix = torch.tensor(
+            [[1, 0, 1, 0], [0, 1, 0, 1], [-0.5, 0, 0.5, 0], [0, -0.5, 0, 0.5]], dtype=boxes.dtype, device=boxes.device
+        )
+        boxes @= convert_matrix
         objectness = x[:, :, 4:5]
-        scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
+        scores, labels = torch.max(x[:, :, 5:], dim=-1, keepdim=True)
         scores *= objectness
-        classes = classes.float()
-        return boxes, scores, classes
-
-
-def suppress_warnings():
-    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
-    warnings.filterwarnings('ignore', category=UserWarning)
-    warnings.filterwarnings('ignore', category=DeprecationWarning)
+        return torch.cat([boxes, scores, labels.to(boxes.dtype)], dim=-1)
 
 
 def yolox_export(weights, exp_file):
@@ -42,10 +37,19 @@ def yolox_export(weights, exp_file):
     return model, exp
 
 
+def suppress_warnings():
+    import warnings
+    warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
+    warnings.filterwarnings('ignore', category=UserWarning)
+    warnings.filterwarnings('ignore', category=DeprecationWarning)
+    warnings.filterwarnings('ignore', category=FutureWarning)
+    warnings.filterwarnings('ignore', category=ResourceWarning)
+
+
 def main(args):
     suppress_warnings()
 
-    print('\nStarting: %s' % args.weights)
+    print(f'\nStarting: {args.weights}')
 
     print('Opening YOLOX model')
 
@@ -57,39 +61,35 @@ def main(args):
     img_size = [exp.input_size[1], exp.input_size[0]]
 
     onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
-    onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx'
+    onnx_output_file = f'{args.weights}.onnx'
 
     dynamic_axes = {
         'input': {
             0: 'batch'
         },
-        'boxes': {
-            0: 'batch'
-        },
-        'scores': {
-            0: 'batch'
-        },
-        'classes': {
+        'output': {
             0: 'batch'
         }
     }
 
     print('Exporting the model to ONNX')
-    torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
-                      do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
-                      dynamic_axes=dynamic_axes if args.dynamic else None)
+    torch.onnx.export(
+        model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset, do_constant_folding=True,
+        input_names=['input'], output_names=['output'], dynamic_axes=dynamic_axes if args.dynamic else None
+    )
 
     if args.simplify:
         print('Simplifying the ONNX model')
-        import onnxsim
+        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('Done: %s\n' % onnx_output_file)
+    print(f'Done: {onnx_output_file}\n')
 
 
 def parse_args():
+    import argparse
     parser = argparse.ArgumentParser(description='DeepStream YOLOX conversion')
     parser.add_argument('-w', '--weights', required=True, help='Input weights (.pth) file path (required)')
     parser.add_argument('-c', '--exp', required=True, help='Input exp (.py) file path (required)')
@@ -100,8 +100,6 @@ def parse_args():
     args = parser.parse_args()
     if not os.path.isfile(args.weights):
         raise SystemExit('Invalid weights file')
-    if not os.path.isfile(args.exp):
-        raise SystemExit('Invalid exp file')
     if args.dynamic and args.batch > 1:
         raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
     return args
@@ -109,4 +107,4 @@ def parse_args():
 
 if __name__ == '__main__':
     args = parse_args()
-    sys.exit(main(args))
+    main(args)