Fix README.md and replace max-wh flag by ort flag

deploy/ONNX/README.md

@@ -26,8 +26,8 @@ python ./deploy/ONNX/export_onnx.py \
 - `--simplify` : Whether to simplify onnx. Not support in end to end export.
 - `--end2end` : Whether to export end to end onnx model. Only support onnxruntime and TensorRT >= 8.0.0 .
 - `--trt-version` :  Export onnx for TensorRT version. Support : 7 or 8.
+- `--ort` : Whether to export onnx for onnxruntime backend.
 - `--with-preprocess` : Whether to export preprocess with bgr2rgb and normalize (divide by 255)
-- `--max-wh` : Default is None for TensorRT backend. Set int for onnxruntime backend.
 - `--topk-all` : Topk objects for every image.
 - `--iou-thres` : IoU threshold for NMS algorithm.
 - `--conf-thres` : Confidence threshold for NMS algorithm.
@@ -58,17 +58,11 @@ python ./deploy/ONNX/export_onnx.py \
     --img 640 \
     --batch 1 \
     --end2end \
-    --max-wh 7680
+    --ort
 ```
 
 You will get an onnx with **NonMaxSuppression** operater .
 
-The onnx outputs shape is ```nums x 7```.
-
-```nums``` means the number of all objects which were detected.
-
-```7```  means [`batch_index`,`x0`,`y0`,`x1`,`y1`,`classid`,`score`]
-
 ### TensorRT backend (TensorRT version == 7.2.3.4)
 #### Usage
 ```bash

deploy/ONNX/export_onnx.py

@@ -31,8 +31,8 @@
     parser.add_argument('--dynamic-batch', action='store_true', help='export dynamic batch onnx model')
     parser.add_argument('--end2end', action='store_true', help='export end2end onnx')
     parser.add_argument('--trt-version', type=int, default=8, help='tensorrt version')
+    parser.add_argument('--ort', action='store_true', help='export onnx for onnxruntime')
     parser.add_argument('--with-preprocess', action='store_true', help='export bgr2rgb and normalize')
-    parser.add_argument('--max-wh', type=int, default=None, help='None for tensorrt nms, int value for onnx-runtime nms')
     parser.add_argument('--topk-all', type=int, default=100, help='topk objects for every images')
     parser.add_argument('--iou-thres', type=float, default=0.45, help='iou threshold for NMS')
     parser.add_argument('--conf-thres', type=float, default=0.4, help='conf threshold for NMS')
@@ -71,7 +71,7 @@
             'images' :{
                 0:'batch',
             },}
-        if args.end2end and args.max_wh is None:
+        if args.end2end and not args.ort:
             output_axes = {
                 'num_dets': {0: 'batch'},
                 'det_boxes': {0: 'batch'},
@@ -88,7 +88,7 @@
     if args.end2end:
         from yolov6.models.end2end import End2End
         model = End2End(model, max_obj=args.topk_all, iou_thres=args.iou_thres,score_thres=args.conf_thres,
-                        max_wh=args.max_wh, device=device, trt_version=args.trt_version, with_preprocess=args.with_preprocess)
+                        device=device, ort=args.ort, trt_version=args.trt_version, with_preprocess=args.with_preprocess)
 
     print("===================")
     print(model)
@@ -113,7 +113,7 @@
             onnx_model = onnx.load(f)  # load onnx model
             onnx.checker.check_model(onnx_model)  # check onnx model
             # Fix output shape
-            if args.end2end and args.max_wh is None:
+            if args.end2end and not args.ort:
                 shapes = [args.batch_size, 1, args.batch_size, args.topk_all, 4,
                           args.batch_size, args.topk_all, args.batch_size, args.topk_all]
                 for i in onnx_model.graph.output:
@@ -135,7 +135,7 @@
     # Finish
     LOGGER.info('\nExport complete (%.2fs)' % (time.time() - t))
     if args.end2end:
-        if args.max_wh is None:
+        if not args.ort:
             info = f'trtexec --onnx={export_file} --saveEngine={export_file.replace(".onnx",".engine")}'
             if args.dynamic_batch:
                 LOGGER.info('Dynamic batch export should define min/opt/max batchsize\n'+

yolov6/models/end2end.py

@@ -139,13 +139,12 @@ def symbolic(g,
 
 class ONNX_ORT(nn.Module):
     '''onnx module with ONNX-Runtime NMS operation.'''
-    def __init__(self, max_obj=100, iou_thres=0.45, score_thres=0.25, max_wh=640, device=None):
+    def __init__(self, max_obj=100, iou_thres=0.45, score_thres=0.25, device=None):
         super().__init__()
         self.device = device if device else torch.device("cpu")
         self.max_obj = torch.tensor([max_obj]).to(device)
         self.iou_threshold = torch.tensor([iou_thres]).to(device)
         self.score_threshold = torch.tensor([score_thres]).to(device)
-        self.max_wh = max_wh
         self.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=torch.float32,
                                            device=self.device)
@@ -191,9 +190,8 @@ def forward(self, x):
 
 class ONNX_TRT7(nn.Module):
     '''onnx module with TensorRT NMS operation.'''
-    def __init__(self, max_obj=100, iou_thres=0.45, score_thres=0.25, max_wh=None ,device=None):
+    def __init__(self, max_obj=100, iou_thres=0.45, score_thres=0.25, device=None):
         super().__init__()
-        assert max_wh is None
         self.device = device if device else torch.device('cpu')
         self.shareLocation = 1
         self.backgroundLabelId = -1
@@ -236,9 +234,8 @@ def forward(self, x):
 
 class ONNX_TRT8(nn.Module):
     '''onnx module with TensorRT NMS operation.'''
-    def __init__(self, max_obj=100, iou_thres=0.45, score_thres=0.25, max_wh=None ,device=None):
+    def __init__(self, max_obj=100, iou_thres=0.45, score_thres=0.25, device=None):
         super().__init__()
-        assert max_wh is None
         self.device = device if device else torch.device('cpu')
         self.background_class = -1,
         self.box_coding = 1,
@@ -262,14 +259,14 @@ def forward(self, x):
 
 class End2End(nn.Module):
     '''export onnx or tensorrt model with NMS operation.'''
-    def __init__(self, model, max_obj=100, iou_thres=0.45, score_thres=0.25, max_wh=None, device=None, trt_version=8, with_preprocess=False):
+    def __init__(self, model, max_obj=100, iou_thres=0.45, score_thres=0.25, device=None, ort=False,  trt_version=8, with_preprocess=False):
         super().__init__()
         device = device if device else torch.device('cpu')
         self.with_preprocess = with_preprocess
         self.model = model.to(device)
         TRT = ONNX_TRT8 if trt_version >= 8  else ONNX_TRT7
-        self.patch_model = TRT if max_wh is None else ONNX_ORT
-        self.end2end = self.patch_model(max_obj, iou_thres, score_thres, max_wh, device)
+        self.patch_model = ONNX_ORT if ort else TRT
+        self.end2end = self.patch_model(max_obj, iou_thres, score_thres, device)
         self.end2end.eval()
 
     def forward(self, x):