thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add YOLO-Face + Minor fixes

Browse Source
This commit is contained in:
Marcos Luciano
2023-09-09 20:32:25 -03:00
parent 1a9df997a4
commit 1445990136
35 changed files with 54 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
README.md
View File

@@ -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-Pose: https://github.com/marcoslucianops/DeepStream-Yolo-Pose
### YOLO-Seg: https://github.com/marcoslucianops/DeepStream-Yolo-Seg ### 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 ### 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 * DeepStream tutorials
* Updated INT8 calibration * Updated INT8 calibration
* Support for segmentation models
* Support for classification models * Support for classification models
### Improvements on this repository ### 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 * 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-Pose: https://github.com/marcoslucianops/DeepStream-Yolo-Pose**
* **YOLO-Seg: https://github.com/marcoslucianops/DeepStream-Yolo-Seg** * **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 #### 5. Run
``` ```

2
config_infer_primary.txt
View File

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

3
config_infer_primary_damoyolo.txt
View File

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

3
config_infer_primary_ppyoloe.txt
View File

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

3
config_infer_primary_ppyoloe_plus.txt
View File

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

3
config_infer_primary_yoloV2.txt
View File

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

3
config_infer_primary_yoloV5.txt
View File

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

3
config_infer_primary_yoloV6.txt
View File

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

3
config_infer_primary_yoloV7.txt
View File

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

3
config_infer_primary_yoloV8.txt
View File

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

3
config_infer_primary_yolonas.txt
View File

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

3
config_infer_primary_yolonas_custom.txt
View File

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

3
config_infer_primary_yolor.txt
View File

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

3
config_infer_primary_yolox.txt
View File

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

3
config_infer_primary_yolox_legacy.txt
View File

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

11
docs/DAMOYOLO.md
View File

@@ -76,7 +76,7 @@ or
--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 --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 ### Edit the deepstream_app_config file

13
docs/PPYOLOE.md
View File

@@ -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) 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 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 --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 --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 ### Edit the deepstream_app_config file

11
docs/YOLONAS.md
View File

@@ -109,7 +109,7 @@ or
--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 --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 ### Edit the deepstream_app_config file

11
docs/YOLOR.md
View File

@@ -94,7 +94,7 @@ or
--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 --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 ### Edit the deepstream_app_config file

11
docs/YOLOX.md
View File

@@ -58,7 +58,7 @@ python3 export_yolox.py -w yolox_s.pth -c exps/default/yolox_s.py --dynamic
--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 --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 ### Edit the deepstream_app_config file

11
docs/YOLOv5.md
View File

@@ -86,7 +86,7 @@ or
--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 --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 ### Edit the deepstream_app_config file

11
docs/YOLOv6.md
View File

@@ -86,7 +86,7 @@ or
--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 --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 ### Edit the deepstream_app_config file

11
docs/YOLOv7.md
View File

@@ -88,7 +88,7 @@ or
--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 --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 ### Edit the deepstream_app_config file

11
docs/YOLOv8.md
View File

@@ -79,7 +79,7 @@ or
--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 --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 ### Edit the deepstream_app_config file

2
nvdsinfer_custom_impl_Yolo/yolo.cpp
View File

@@ -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(); nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile();
assert(profile); assert(profile);
for (INT i = 0; i < network->getNbInputs(); ++i) { for (INT i = 0; i < network->getNbInputs(); ++i) {

4
utils/export_damoyolo.py
View File

@@ -107,14 +107,14 @@ def parse_args():
parser.add_argument('--opset', type=int, default=11, help='ONNX opset version') 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('--simplify', action='store_true', help='ONNX simplify model')
parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size') 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() args = parser.parse_args()
if not os.path.isfile(args.weights): if not os.path.isfile(args.weights):
raise SystemExit('Invalid weights file') raise SystemExit('Invalid weights file')
if not os.path.isfile(args.config): if not os.path.isfile(args.config):
raise SystemExit('Invalid config file') raise SystemExit('Invalid config file')
if args.dynamic and args.batch > 1: 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 return args

4
utils/export_ppyoloe.py
View File

@@ -89,12 +89,12 @@ def parse_args():
parser.add_argument('--opset', type=int, default=11, help='ONNX opset version') 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('--simplify', action='store_true', help='ONNX simplify model')
parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size') 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() args = parser.parse_args()
if not os.path.isfile(args.weights): if not os.path.isfile(args.weights):
raise SystemExit('\nInvalid weights file') raise SystemExit('\nInvalid weights file')
if args.dynamic and args.batch > 1: 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: elif args.dynamic:
args.batch = None args.batch = None
return args return args

6
utils/export_yoloV5.py
View File

@@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module):
x = x[0] x = x[0]
boxes = x[:, :, :4] boxes = x[:, :, :4]
objectness = x[:, :, 4:5] 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 scores *= objectness
classes = classes.float() classes = classes.float()
return boxes, scores, classes 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('--opset', type=int, default=17, help='ONNX opset version')
parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') 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('--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() args = parser.parse_args()
if not os.path.isfile(args.weights): if not os.path.isfile(args.weights):
raise SystemExit('Invalid weights file') raise SystemExit('Invalid weights file')
if args.dynamic and args.batch > 1: 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 return args

6
utils/export_yoloV6.py
View File

@@ -22,7 +22,7 @@ class DeepStreamOutput(nn.Module):
def forward(self, x): def forward(self, x):
boxes = x[:, :, :4] boxes = x[:, :, :4]
objectness = x[:, :, 4:5] 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 scores *= objectness
classes = classes.float() classes = classes.float()
return boxes, scores, classes 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('--opset', type=int, default=13, help='ONNX opset version')
parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') 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('--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() args = parser.parse_args()
if not os.path.isfile(args.weights): if not os.path.isfile(args.weights):
raise SystemExit('Invalid weights file') raise SystemExit('Invalid weights file')
if args.dynamic and args.batch > 1: 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 return args

6
utils/export_yoloV7.py
View File

@@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module):
def forward(self, x): def forward(self, x):
boxes = x[:, :, :4] boxes = x[:, :, :4]
objectness = x[:, :, 4:5] 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 scores *= objectness
classes = classes.float() classes = classes.float()
return boxes, scores, classes 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('--opset', type=int, default=12, help='ONNX opset version')
parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') 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('--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() args = parser.parse_args()
if not os.path.isfile(args.weights): if not os.path.isfile(args.weights):
raise SystemExit('Invalid weights file') raise SystemExit('Invalid weights file')
if args.dynamic and args.batch > 1: 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 return args

6
utils/export_yoloV7_u6.py
View File

@@ -17,7 +17,7 @@ class DeepStreamOutput(nn.Module):
def forward(self, x): def forward(self, x):
x = x.transpose(1, 2) x = x.transpose(1, 2)
boxes = x[:, :, :4] 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() classes = classes.float()
return boxes, scores, classes 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('--opset', type=int, default=12, help='ONNX opset version')
parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') 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('--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() args = parser.parse_args()
if not os.path.isfile(args.weights): if not os.path.isfile(args.weights):
raise SystemExit('Invalid weights file') raise SystemExit('Invalid weights file')
if args.dynamic and args.batch > 1: 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 return args

6
utils/export_yoloV8.py
View File

@@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module):
def forward(self, x): def forward(self, x):
x = x.transpose(1, 2) x = x.transpose(1, 2)
boxes = x[:, :, :4] 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() classes = classes.float()
return boxes, scores, classes 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('--opset', type=int, default=16, help='ONNX opset version')
parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') 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('--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() args = parser.parse_args()
if not os.path.isfile(args.weights): if not os.path.isfile(args.weights):
raise SystemExit('Invalid weights file') raise SystemExit('Invalid weights file')
if args.dynamic and args.batch > 1: 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 return args

4
utils/export_yolonas.py
View File

@@ -89,14 +89,14 @@ def parse_args():
parser.add_argument('--opset', type=int, default=14, help='ONNX opset version') 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('--simplify', action='store_true', help='ONNX simplify model')
parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size') 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() args = parser.parse_args()
if args.model == '': if args.model == '':
raise SystemExit('Invalid model name') raise SystemExit('Invalid model name')
if not os.path.isfile(args.weights): if not os.path.isfile(args.weights):
raise SystemExit('Invalid weights file') raise SystemExit('Invalid weights file')
if args.dynamic and args.batch > 1: 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 return args

6
utils/export_yolor.py
View File

@@ -15,7 +15,7 @@ class DeepStreamOutput(nn.Module):
x = x[0] x = x[0]
boxes = x[:, :, :4] boxes = x[:, :, :4]
objectness = x[:, :, 4:5] 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 scores *= objectness
classes = classes.float() classes = classes.float()
return boxes, scores, classes 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('--opset', type=int, default=12, help='ONNX opset version')
parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') 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('--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() args = parser.parse_args()
if not os.path.isfile(args.weights): if not os.path.isfile(args.weights):
raise SystemExit('Invalid weights file') raise SystemExit('Invalid weights file')
if args.dynamic and args.batch > 1: 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 return args

6
utils/export_yolox.py
View File

@@ -17,7 +17,7 @@ class DeepStreamOutput(nn.Module):
def forward(self, x): def forward(self, x):
boxes = x[:, :, :4] boxes = x[:, :, :4]
objectness = x[:, :, 4:5] 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 scores *= objectness
classes = classes.float() classes = classes.float()
return boxes, scores, classes 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('--opset', type=int, default=11, help='ONNX opset version')
parser.add_argument('--simplify', action='store_true', help='ONNX simplify model') 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('--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() args = parser.parse_args()
if not os.path.isfile(args.weights): if not os.path.isfile(args.weights):
raise SystemExit('Invalid weights file') raise SystemExit('Invalid weights file')
if not os.path.isfile(args.exp): if not os.path.isfile(args.exp):
raise SystemExit('Invalid exp file') raise SystemExit('Invalid exp file')
if args.dynamic and args.batch > 1: 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 return args