Add YOLO-Face + Minor fixes

README.md

@@ -5,6 +5,7 @@ NVIDIA DeepStream SDK 6.3 / 6.2 / 6.1.1 / 6.1 / 6.0.1 / 6.0 / 5.1  configuration
 --------------------------------------------------------------------------------------------------
 ### YOLO-Pose: https://github.com/marcoslucianops/DeepStream-Yolo-Pose
 ### YOLO-Seg: https://github.com/marcoslucianops/DeepStream-Yolo-Seg
+### YOLO-Face: https://github.com/marcoslucianops/DeepStream-Yolo-Face
 --------------------------------------------------------------------------------------------------
 ### Important: please export the ONNX model with the new export file, generate the TensorRT engine again with the updated files, and use the new config_infer_primary file according to your model
 --------------------------------------------------------------------------------------------------
@@ -13,7 +14,6 @@ NVIDIA DeepStream SDK 6.3 / 6.2 / 6.1.1 / 6.1 / 6.0.1 / 6.0 / 5.1  configuration
 
 * DeepStream tutorials
 * Updated INT8 calibration
-* Support for segmentation models
 * Support for classification models
 
 ### Improvements on this repository
@@ -31,6 +31,7 @@ NVIDIA DeepStream SDK 6.3 / 6.2 / 6.1.1 / 6.1 / 6.0.1 / 6.0 / 5.1  configuration
 * New output structure (fix wrong output on DeepStream < 6.2) - it need to export the ONNX model with the new export file, generate the TensorRT engine again with the updated files, and use the new config_infer_primary file according to your model
 * **YOLO-Pose: https://github.com/marcoslucianops/DeepStream-Yolo-Pose**
 * **YOLO-Seg: https://github.com/marcoslucianops/DeepStream-Yolo-Seg**
+* **YOLO-Face: https://github.com/marcoslucianops/DeepStream-Yolo-Face**
 
 ##
 
@@ -246,14 +247,6 @@ model-file=yolov4.weights
 ...
 ```
 
-**NOTE**: By default, the dynamic batch-size is set. To use implicit batch-size, uncomment the line
-
-```
-...
-force-implicit-batch-dim=1
-...
-```
-
 #### 5. Run
 
 ```

config_infer_primary.txt

@@ -18,7 +18,7 @@ cluster-mode=2
 maintain-aspect-ratio=0
 symmetric-padding=1
 #force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYolo
 #parse-bbox-func-name=NvDsInferParseYoloCuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_damoyolo.txt

@@ -15,8 +15,7 @@ process-mode=1
 network-type=0
 cluster-mode=2
 maintain-aspect-ratio=0
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYoloE
 #parse-bbox-func-name=NvDsInferParseYoloECuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_ppyoloe.txt

@@ -16,8 +16,7 @@ process-mode=1
 network-type=0
 cluster-mode=2
 maintain-aspect-ratio=0
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYoloE
 #parse-bbox-func-name=NvDsInferParseYoloECuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_ppyoloe_plus.txt

@@ -15,8 +15,7 @@ process-mode=1
 network-type=0
 cluster-mode=2
 maintain-aspect-ratio=0
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYoloE
 #parse-bbox-func-name=NvDsInferParseYoloECuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_yoloV2.txt

@@ -16,8 +16,7 @@ process-mode=1
 network-type=0
 cluster-mode=2
 maintain-aspect-ratio=0
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYolo
 #parse-bbox-func-name=NvDsInferParseYoloCuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_yoloV5.txt

@@ -16,8 +16,7 @@ network-type=0
 cluster-mode=2
 maintain-aspect-ratio=1
 symmetric-padding=1
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYolo
 #parse-bbox-func-name=NvDsInferParseYoloCuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_yoloV6.txt

@@ -16,8 +16,7 @@ network-type=0
 cluster-mode=2
 maintain-aspect-ratio=1
 symmetric-padding=1
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYolo
 #parse-bbox-func-name=NvDsInferParseYoloCuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_yoloV7.txt

@@ -16,8 +16,7 @@ network-type=0
 cluster-mode=2
 maintain-aspect-ratio=1
 symmetric-padding=1
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYolo
 #parse-bbox-func-name=NvDsInferParseYoloCuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_yoloV8.txt

@@ -16,8 +16,7 @@ network-type=0
 cluster-mode=2
 maintain-aspect-ratio=1
 symmetric-padding=1
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYolo
 #parse-bbox-func-name=NvDsInferParseYoloCuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_yolonas.txt

@@ -16,8 +16,7 @@ network-type=0
 cluster-mode=2
 maintain-aspect-ratio=1
 symmetric-padding=0
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYoloE
 #parse-bbox-func-name=NvDsInferParseYoloECuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_yolonas_custom.txt

@@ -16,8 +16,7 @@ network-type=0
 cluster-mode=2
 maintain-aspect-ratio=1
 symmetric-padding=0
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYoloE
 #parse-bbox-func-name=NvDsInferParseYoloECuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_yolor.txt

@@ -16,8 +16,7 @@ network-type=0
 cluster-mode=2
 maintain-aspect-ratio=1
 symmetric-padding=1
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYolo
 #parse-bbox-func-name=NvDsInferParseYoloCuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_yolox.txt

@@ -16,8 +16,7 @@ network-type=0
 cluster-mode=2
 maintain-aspect-ratio=1
 symmetric-padding=0
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYolo
 #parse-bbox-func-name=NvDsInferParseYoloCuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

config_infer_primary_yolox_legacy.txt

@@ -17,8 +17,7 @@ network-type=0
 cluster-mode=2
 maintain-aspect-ratio=1
 symmetric-padding=0
-#force-implicit-batch-dim=1
-#workspace-size=1000
+#workspace-size=2000
 parse-bbox-func-name=NvDsInferParseYolo
 #parse-bbox-func-name=NvDsInferParseYoloCuda
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so

docs/DAMOYOLO.md

@@ -76,7 +76,7 @@ or
 --dynamic
 ```
 
-**NOTE**: To use implicit batch-size (example for batch-size = 4)
+**NOTE**: To use static batch-size (example for batch-size = 4)
 
 ```
 --batch 4
@@ -172,15 +172,6 @@ maintain-aspect-ratio=0
 ...
 ```
 
-**NOTE**: By default, the dynamic batch-size is set. To use implicit batch-size, uncomment the line
-
-```
-[property]
-...
-force-implicit-batch-dim=1
-...
-```
-
 ##
 
 ### Edit the deepstream_app_config file

docs/PPYOLOE.md

@@ -37,7 +37,7 @@ wget https://paddledet.bj.bcebos.com/models/ppyoloe_plus_crn_s_80e_coco.pdparams
 Generate the ONNX model file (example for PP-YOLOE+_s)
 
 ```
-pip3 install onnx onnxsim onnxruntime
+pip3 install onnx onnxsim onnxruntime paddle2onnx
 python3 export_ppyoloe.py -w ppyoloe_plus_crn_s_80e_coco.pdparams -c configs/ppyoloe/ppyoloe_plus_crn_s_80e_coco.yml --dynamic
 ```
 
@@ -53,7 +53,7 @@ python3 export_ppyoloe.py -w ppyoloe_plus_crn_s_80e_coco.pdparams -c configs/ppy
 --dynamic
 ```
 
-**NOTE**: To use implicit batch-size (example for batch-size = 4)
+**NOTE**: To use static batch-size (example for batch-size = 4)
 
 ```
 --batch 4
@@ -172,15 +172,6 @@ offsets=123.675;116.28;103.53
 ...
 ```
 
-**NOTE**: By default, the dynamic batch-size is set. To use implicit batch-size, uncomment the line
-
-```
-[property]
-...
-force-implicit-batch-dim=1
-...
-```
-
 ##
 
 ### Edit the deepstream_app_config file

docs/YOLONAS.md

@@ -109,7 +109,7 @@ or
 --dynamic
 ```
 
-**NOTE**: To use implicit batch-size (example for batch-size = 4)
+**NOTE**: To use static batch-size (example for batch-size = 4)
 
 ```
 --batch 4
@@ -226,15 +226,6 @@ net-scale-factor=1
 ...
 ```
 
-**NOTE**: By default, the dynamic batch-size is set. To use implicit batch-size, uncomment the line
-
-```
-[property]
-...
-force-implicit-batch-dim=1
-...
-```
-
 ##
 
 ### Edit the deepstream_app_config file

docs/YOLOR.md

@@ -94,7 +94,7 @@ or
 --dynamic
 ```
 
-**NOTE**: To use implicit batch-size (example for batch-size = 4)
+**NOTE**: To use static batch-size (example for batch-size = 4)
 
 ```
 --batch 4
@@ -191,15 +191,6 @@ symmetric-padding=1
 ...
 ```
 
-**NOTE**: By default, the dynamic batch-size is set. To use implicit batch-size, uncomment the line
-
-```
-[property]
-...
-force-implicit-batch-dim=1
-...
-```
-
 ##
 
 ### Edit the deepstream_app_config file

docs/YOLOX.md

@@ -58,7 +58,7 @@ python3 export_yolox.py -w yolox_s.pth -c exps/default/yolox_s.py --dynamic
 --dynamic
 ```
 
-**NOTE**: To use implicit batch-size (example for batch-size = 4)
+**NOTE**: To use static batch-size (example for batch-size = 4)
 
 ```
 --batch 4
@@ -178,15 +178,6 @@ offsets=123.675;116.28;103.53
 ...
 ```
 
-**NOTE**: By default, the dynamic batch-size is set. To use implicit batch-size, uncomment the line
-
-```
-[property]
-...
-force-implicit-batch-dim=1
-...
-```
-
 ##
 
 ### Edit the deepstream_app_config file

docs/YOLOv5.md

@@ -86,7 +86,7 @@ or
 --dynamic
 ```
 
-**NOTE**: To use implicit batch-size (example for batch-size = 4)
+**NOTE**: To use static batch-size (example for batch-size = 4)
 
 ```
 --batch 4
@@ -183,15 +183,6 @@ symmetric-padding=1
 ...
 ```
 
-**NOTE**: By default, the dynamic batch-size is set. To use implicit batch-size, uncomment the line
-
-```
-[property]
-...
-force-implicit-batch-dim=1
-...
-```
-
 ##
 
 ### Edit the deepstream_app_config file

docs/YOLOv6.md

@@ -86,7 +86,7 @@ or
 --dynamic
 ```
 
-**NOTE**: To use implicit batch-size (example for batch-size = 4)
+**NOTE**: To use static batch-size (example for batch-size = 4)
 
 ```
 --batch 4
@@ -183,15 +183,6 @@ symmetric-padding=1
 ...
 ```
 
-**NOTE**: By default, the dynamic batch-size is set. To use implicit batch-size, uncomment the line
-
-```
-[property]
-...
-force-implicit-batch-dim=1
-...
-```
-
 ##
 
 ### Edit the deepstream_app_config file

docs/YOLOv7.md

@@ -88,7 +88,7 @@ or
 --dynamic
 ```
 
-**NOTE**: To use implicit batch-size (example for batch-size = 4)
+**NOTE**: To use static batch-size (example for batch-size = 4)
 
 ```
 --batch 4
@@ -185,15 +185,6 @@ symmetric-padding=1
 ...
 ```
 
-**NOTE**: By default, the dynamic batch-size is set. To use implicit batch-size, uncomment the line
-
-```
-[property]
-...
-force-implicit-batch-dim=1
-...
-```
-
 ##
 
 ### Edit the deepstream_app_config file

docs/YOLOv8.md

@@ -79,7 +79,7 @@ or
 --dynamic
 ```
 
-**NOTE**: To use implicit batch-size (example for batch-size = 4)
+**NOTE**: To use static batch-size (example for batch-size = 4)
 
 ```
 --batch 4
@@ -176,15 +176,6 @@ symmetric-padding=1
 ...
 ```
 
-**NOTE**: By default, the dynamic batch-size is set. To use implicit batch-size, uncomment the line
-
-```
-[property]
-...
-force-implicit-batch-dim=1
-...
-```
-
 ##
 
 ### Edit the deepstream_app_config file

nvdsinfer_custom_impl_Yolo/yolo.cpp

@@ -116,7 +116,7 @@ Yolo::createEngine(nvinfer1::IBuilder* builder)
     }
   }
 
-  if (!m_ImplicitBatch && network->getInput(0)->getDimensions().d[0] == -1) {
+  if ((m_NetworkType == "darknet" && !m_ImplicitBatch) || network->getInput(0)->getDimensions().d[0] == -1) {
     nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile();
     assert(profile);
     for (INT i = 0; i < network->getNbInputs(); ++i) {

utils/export_damoyolo.py

@@ -107,14 +107,14 @@ def parse_args():
     parser.add_argument('--opset', type=int, default=11, 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='Implicit 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 implicit batch-size at same time')
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
     return args
 
 

utils/export_ppyoloe.py

@@ -89,12 +89,12 @@ def parse_args():
     parser.add_argument('--opset', type=int, default=11, 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='Implicit 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('\nInvalid weights file')
     if args.dynamic and args.batch > 1:
-        raise SystemExit('\nCannot set dynamic batch-size and implicit batch-size at same time')
+        raise SystemExit('\nCannot set dynamic batch-size and static batch-size at same time')
     elif args.dynamic:
         args.batch = None
     return args

utils/export_yoloV5.py

@@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module):
         x = x[0]
         boxes = x[:, :, :4]
         objectness = x[:, :, 4:5]
-        scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
+        scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True)
         scores *= objectness
         classes = classes.float()
         return boxes, scores, classes
@@ -106,12 +106,12 @@ def parse_args():
     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='Implicit 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 implicit batch-size at same time')
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
     return args
 
 

utils/export_yoloV6.py

@@ -22,7 +22,7 @@ class DeepStreamOutput(nn.Module):
     def forward(self, x):
         boxes = x[:, :, :4]
         objectness = x[:, :, 4:5]
-        scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
+        scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True)
         scores *= objectness
         classes = classes.float()
         return boxes, scores, classes
@@ -109,12 +109,12 @@ def parse_args():
     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='Implicit 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 implicit batch-size at same time')
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
     return args
 
 

utils/export_yoloV7.py

@@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module):
     def forward(self, x):
         boxes = x[:, :, :4]
         objectness = x[:, :, 4:5]
-        scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
+        scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True)
         scores *= objectness
         classes = classes.float()
         return boxes, scores, classes
@@ -110,12 +110,12 @@ def parse_args():
     parser.add_argument('--opset', type=int, default=12, 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='Implicit 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 implicit batch-size at same time')
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
     return args
 
 

utils/export_yoloV7_u6.py

@@ -17,7 +17,7 @@ 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)
+        scores, classes = torch.max(x[:, :, 4:5], 2, keepdim=True)
         classes = classes.float()
         return boxes, scores, classes
 
@@ -100,12 +100,12 @@ def parse_args():
     parser.add_argument('--opset', type=int, default=12, 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='Implicit 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 implicit batch-size at same time')
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
     return args
 
 

utils/export_yoloV8.py

@@ -18,7 +18,7 @@ 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)
+        scores, classes = torch.max(x[:, :, 4:5], 2, keepdim=True)
         classes = classes.float()
         return boxes, scores, classes
 
@@ -108,12 +108,12 @@ def parse_args():
     parser.add_argument('--opset', type=int, default=16, 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='Implicit 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 implicit batch-size at same time')
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
     return args
 
 

utils/export_yolonas.py

@@ -89,14 +89,14 @@ def parse_args():
     parser.add_argument('--opset', type=int, default=14, 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='Implicit batch-size')
+    parser.add_argument('--batch', type=int, default=1, help='Static batch-size')
     args = parser.parse_args()
     if args.model == '':
         raise SystemExit('Invalid model name')
     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 implicit batch-size at same time')
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
     return args
 
 

utils/export_yolor.py

@@ -15,7 +15,7 @@ class DeepStreamOutput(nn.Module):
         x = x[0]
         boxes = x[:, :, :4]
         objectness = x[:, :, 4:5]
-        scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
+        scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True)
         scores *= objectness
         classes = classes.float()
         return boxes, scores, classes
@@ -123,12 +123,12 @@ def parse_args():
     parser.add_argument('--opset', type=int, default=12, 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='Implicit 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 implicit batch-size at same time')
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
     return args
 
 

utils/export_yolox.py

@@ -17,7 +17,7 @@ class DeepStreamOutput(nn.Module):
     def forward(self, x):
         boxes = x[:, :, :4]
         objectness = x[:, :, 4:5]
-        scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
+        scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True)
         scores *= objectness
         classes = classes.float()
         return boxes, scores, classes
@@ -96,14 +96,14 @@ def parse_args():
     parser.add_argument('--opset', type=int, default=11, 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='Implicit 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.exp):
         raise SystemExit('Invalid exp file')
     if args.dynamic and args.batch > 1:
-        raise SystemExit('Cannot set dynamic batch-size and implicit batch-size at same time')
+        raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
     return args