From 1445990136518da97121fa7df2f2a89cba2051de Mon Sep 17 00:00:00 2001 From: Marcos Luciano Date: Sat, 9 Sep 2023 20:32:25 -0300 Subject: [PATCH] Add YOLO-Face + Minor fixes --- README.md | 11 ++--------- config_infer_primary.txt | 2 +- config_infer_primary_damoyolo.txt | 3 +-- config_infer_primary_ppyoloe.txt | 3 +-- config_infer_primary_ppyoloe_plus.txt | 3 +-- config_infer_primary_yoloV2.txt | 3 +-- config_infer_primary_yoloV5.txt | 3 +-- config_infer_primary_yoloV6.txt | 3 +-- config_infer_primary_yoloV7.txt | 3 +-- config_infer_primary_yoloV8.txt | 3 +-- config_infer_primary_yolonas.txt | 3 +-- config_infer_primary_yolonas_custom.txt | 3 +-- config_infer_primary_yolor.txt | 3 +-- config_infer_primary_yolox.txt | 3 +-- config_infer_primary_yolox_legacy.txt | 3 +-- docs/DAMOYOLO.md | 11 +---------- docs/PPYOLOE.md | 13 ++----------- docs/YOLONAS.md | 11 +---------- docs/YOLOR.md | 11 +---------- docs/YOLOX.md | 11 +---------- docs/YOLOv5.md | 11 +---------- docs/YOLOv6.md | 11 +---------- docs/YOLOv7.md | 11 +---------- docs/YOLOv8.md | 11 +---------- nvdsinfer_custom_impl_Yolo/yolo.cpp | 2 +- utils/export_damoyolo.py | 4 ++-- utils/export_ppyoloe.py | 4 ++-- utils/export_yoloV5.py | 6 +++--- utils/export_yoloV6.py | 6 +++--- utils/export_yoloV7.py | 6 +++--- utils/export_yoloV7_u6.py | 6 +++--- utils/export_yoloV8.py | 6 +++--- utils/export_yolonas.py | 4 ++-- utils/export_yolor.py | 6 +++--- utils/export_yolox.py | 6 +++--- 35 files changed, 54 insertions(+), 155 deletions(-) diff --git a/README.md b/README.md index 3ae4ab2..963d8d1 100644 --- a/README.md +++ b/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 ``` diff --git a/config_infer_primary.txt b/config_infer_primary.txt index db8dc31..041af5a 100644 --- a/config_infer_primary.txt +++ b/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 diff --git a/config_infer_primary_damoyolo.txt b/config_infer_primary_damoyolo.txt index 0befd68..4c55b21 100644 --- a/config_infer_primary_damoyolo.txt +++ b/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 diff --git a/config_infer_primary_ppyoloe.txt b/config_infer_primary_ppyoloe.txt index 512874a..e1f0558 100644 --- a/config_infer_primary_ppyoloe.txt +++ b/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 diff --git a/config_infer_primary_ppyoloe_plus.txt b/config_infer_primary_ppyoloe_plus.txt index 0396752..9180720 100644 --- a/config_infer_primary_ppyoloe_plus.txt +++ b/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 diff --git a/config_infer_primary_yoloV2.txt b/config_infer_primary_yoloV2.txt index cfc313a..690578e 100644 --- a/config_infer_primary_yoloV2.txt +++ b/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 diff --git a/config_infer_primary_yoloV5.txt b/config_infer_primary_yoloV5.txt index 6254e70..4c956e0 100644 --- a/config_infer_primary_yoloV5.txt +++ b/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 diff --git a/config_infer_primary_yoloV6.txt b/config_infer_primary_yoloV6.txt index 27d4ad9..8caf9c6 100644 --- a/config_infer_primary_yoloV6.txt +++ b/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 diff --git a/config_infer_primary_yoloV7.txt b/config_infer_primary_yoloV7.txt index f73d2a5..52bafc6 100644 --- a/config_infer_primary_yoloV7.txt +++ b/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 diff --git a/config_infer_primary_yoloV8.txt b/config_infer_primary_yoloV8.txt index 03fe22b..c0c1311 100644 --- a/config_infer_primary_yoloV8.txt +++ b/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 diff --git a/config_infer_primary_yolonas.txt b/config_infer_primary_yolonas.txt index 80b1610..5f32a87 100644 --- a/config_infer_primary_yolonas.txt +++ b/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 diff --git a/config_infer_primary_yolonas_custom.txt b/config_infer_primary_yolonas_custom.txt index fee3dfb..5c194bc 100644 --- a/config_infer_primary_yolonas_custom.txt +++ b/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 diff --git a/config_infer_primary_yolor.txt b/config_infer_primary_yolor.txt index 468f903..4b28762 100644 --- a/config_infer_primary_yolor.txt +++ b/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 diff --git a/config_infer_primary_yolox.txt b/config_infer_primary_yolox.txt index 5540eba..cbe973a 100644 --- a/config_infer_primary_yolox.txt +++ b/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 diff --git a/config_infer_primary_yolox_legacy.txt b/config_infer_primary_yolox_legacy.txt index 443dd6b..0f73a71 100644 --- a/config_infer_primary_yolox_legacy.txt +++ b/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 diff --git a/docs/DAMOYOLO.md b/docs/DAMOYOLO.md index 3e73201..f243a58 100644 --- a/docs/DAMOYOLO.md +++ b/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 diff --git a/docs/PPYOLOE.md b/docs/PPYOLOE.md index 913468b..9f57f21 100644 --- a/docs/PPYOLOE.md +++ b/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 diff --git a/docs/YOLONAS.md b/docs/YOLONAS.md index 29553b5..121c8cc 100644 --- a/docs/YOLONAS.md +++ b/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 diff --git a/docs/YOLOR.md b/docs/YOLOR.md index 77c8bcb..a85ecd3 100644 --- a/docs/YOLOR.md +++ b/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 diff --git a/docs/YOLOX.md b/docs/YOLOX.md index 43f65ee..7e99240 100644 --- a/docs/YOLOX.md +++ b/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 diff --git a/docs/YOLOv5.md b/docs/YOLOv5.md index e1fed34..44d1313 100644 --- a/docs/YOLOv5.md +++ b/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 diff --git a/docs/YOLOv6.md b/docs/YOLOv6.md index 0c75cec..13d35b5 100644 --- a/docs/YOLOv6.md +++ b/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 diff --git a/docs/YOLOv7.md b/docs/YOLOv7.md index 5e7cda8..cd36b83 100644 --- a/docs/YOLOv7.md +++ b/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 diff --git a/docs/YOLOv8.md b/docs/YOLOv8.md index 68edd1d..5879248 100644 --- a/docs/YOLOv8.md +++ b/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 diff --git a/nvdsinfer_custom_impl_Yolo/yolo.cpp b/nvdsinfer_custom_impl_Yolo/yolo.cpp index affc5ba..1e9023d 100644 --- a/nvdsinfer_custom_impl_Yolo/yolo.cpp +++ b/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) { diff --git a/utils/export_damoyolo.py b/utils/export_damoyolo.py index 5bc26a6..6444199 100644 --- a/utils/export_damoyolo.py +++ b/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 diff --git a/utils/export_ppyoloe.py b/utils/export_ppyoloe.py index d8c065a..6207e44 100644 --- a/utils/export_ppyoloe.py +++ b/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 diff --git a/utils/export_yoloV5.py b/utils/export_yoloV5.py index f752b97..7e7b0f5 100644 --- a/utils/export_yoloV5.py +++ b/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 diff --git a/utils/export_yoloV6.py b/utils/export_yoloV6.py index 912a4fe..abe09f9 100644 --- a/utils/export_yoloV6.py +++ b/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 diff --git a/utils/export_yoloV7.py b/utils/export_yoloV7.py index 718c762..b87a030 100644 --- a/utils/export_yoloV7.py +++ b/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 diff --git a/utils/export_yoloV7_u6.py b/utils/export_yoloV7_u6.py index 25f34ca..0da10c2 100644 --- a/utils/export_yoloV7_u6.py +++ b/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 diff --git a/utils/export_yoloV8.py b/utils/export_yoloV8.py index f996016..df7f0f7 100644 --- a/utils/export_yoloV8.py +++ b/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 diff --git a/utils/export_yolonas.py b/utils/export_yolonas.py index 71356f2..ae80e61 100644 --- a/utils/export_yolonas.py +++ b/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 diff --git a/utils/export_yolor.py b/utils/export_yolor.py index 9a4fad2..e03a3e4 100644 --- a/utils/export_yolor.py +++ b/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 diff --git a/utils/export_yolox.py b/utils/export_yolox.py index 8426381..2b52113 100644 --- a/utils/export_yolox.py +++ b/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