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

Add RT-DETR Ultralytics

Browse Source
This commit is contained in:
Marcos Luciano
2023-11-23 21:08:37 -03:00
parent 5af9da189d
commit 758b7a0bb7
4 changed files with 327 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@@ -30,6 +30,7 @@ NVIDIA DeepStream SDK 6.3 / 6.2 / 6.1.1 / 6.1 / 6.0.1 / 6.0 / 5.1 configuration
* INT8 calibration (PTQ) for Darknet and ONNX exported models
* 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
* **RT-DETR (https://github.com/lyuwenyu/RT-DETR/tree/main/rtdetr_pytorch)**
* **RT-DETR Ultralytics (https://docs.ultralytics.com/models/rtdetr)**
##
@@ -53,6 +54,7 @@ NVIDIA DeepStream SDK 6.3 / 6.2 / 6.1.1 / 6.1 / 6.0.1 / 6.0 / 5.1 configuration
* [PP-YOLOE / PP-YOLOE+ usage](docs/PPYOLOE.md)
* [YOLO-NAS usage](docs/YOLONAS.md)
* [RT-DETR usage](docs/RTDETR.md)
* [RT-DETR Ultralytics usage](docs/RTDETR_Ultralytics.md)
* [Using your custom model](docs/customModels.md)
* [Multiple YOLO GIEs](docs/multipleGIEs.md)

4
docs/RTDETR.md
View File

@@ -1,6 +1,6 @@
# RT_DETR usage
# RT-DETR usage
**NOTE**: For it is supported only the https://github.com/lyuwenyu/RT-DETR/tree/main/rtdetr_pytorch version.
**NOTE**: https://github.com/lyuwenyu/RT-DETR/tree/main/rtdetr_pytorch version.
* [Convert model](#convert-model)
* [Compile the lib](#compile-the-lib)

199
docs/RTDETR_Ultralytics.md Normal file
View File

@@ -0,0 +1,199 @@
# RT-DETR Ultralytics usage
**NOTE**: Ultralytics (https://docs.ultralytics.com/models/rtdetr) version.
* [Convert model](#convert-model)
* [Compile the lib](#compile-the-lib)
* [Edit the config_infer_primary_rtdetr file](#edit-the-config_infer_primary_rtdetr-file)
* [Edit the deepstream_app_config file](#edit-the-deepstream_app_config-file)
* [Testing the model](#testing-the-model)
##
### Convert model
#### 1. Download the Ultralytics repo and install the requirements
```
git clone https://github.com/ultralytics/ultralytics.git
cd ultralytics
pip3 install -r requirements.txt
python3 setup.py install
pip3 install onnx onnxsim onnxruntime
```
**NOTE**: It is recommended to use Python virtualenv.
#### 2. Copy conversor
Copy the `export_rtdetr_ultralytics.py` file from `DeepStream-Yolo/utils` directory to the `ultralytics` folder.
#### 3. Download the model
Download the `pt` file from [Ultralytics](https://github.com/ultralytics/assets/releases/) releases (example for RT-DETR-l)
```
wget https://github.com/ultralytics/assets/releases/download/v0.0.0/rtdetr-l.pt
```
**NOTE**: You can use your custom model.
#### 4. Convert model
Generate the ONNX model file (example for RT-DETR-l)
```
python3 export_rtdetr_ultralytics.py -w rtdetr-l.pt --dynamic
```
**NOTE**: To change the inference size (defaut: 640)
```
-s SIZE
--size SIZE
-s HEIGHT WIDTH
--size HEIGHT WIDTH
```
Example for 1280
```
-s 1280
```
or
```
-s 1280 1280
```
**NOTE**: To simplify the ONNX model (DeepStream >= 6.0)
```
--simplify
```
**NOTE**: To use dynamic batch-size (DeepStream >= 6.1)
```
--dynamic
```
**NOTE**: To use static batch-size (example for batch-size = 4)
```
--batch 4
```
**NOTE**: If you are using the DeepStream 5.1, remove the `--dynamic` arg and use opset 12 or lower. The default opset is 16.
```
--opset 12
```
#### 5. Copy generated files
Copy the generated ONNX model file and labels.txt file (if generated) to the `DeepStream-Yolo` folder.
##
### Compile the lib
Open the `DeepStream-Yolo` folder and compile the lib
* DeepStream 6.3 on x86 platform
```
CUDA_VER=12.1 make -C nvdsinfer_custom_impl_Yolo
```
* DeepStream 6.2 on x86 platform
```
CUDA_VER=11.8 make -C nvdsinfer_custom_impl_Yolo
```
* DeepStream 6.1.1 on x86 platform
```
CUDA_VER=11.7 make -C nvdsinfer_custom_impl_Yolo
```
* DeepStream 6.1 on x86 platform
```
CUDA_VER=11.6 make -C nvdsinfer_custom_impl_Yolo
```
* DeepStream 6.0.1 / 6.0 on x86 platform
```
CUDA_VER=11.4 make -C nvdsinfer_custom_impl_Yolo
```
* DeepStream 5.1 on x86 platform
```
CUDA_VER=11.1 make -C nvdsinfer_custom_impl_Yolo
```
* DeepStream 6.3 / 6.2 / 6.1.1 / 6.1 on Jetson platform
```
CUDA_VER=11.4 make -C nvdsinfer_custom_impl_Yolo
```
* DeepStream 6.0.1 / 6.0 / 5.1 on Jetson platform
```
CUDA_VER=10.2 make -C nvdsinfer_custom_impl_Yolo
```
##
### Edit the config_infer_primary_rtdetr file
Edit the `config_infer_primary_rtdetr.txt` file according to your model (example for RT-DETR-l with 80 classes)
```
[property]
...
onnx-file=rtdetr-l.onnx
...
num-detected-classes=80
...
parse-bbox-func-name=NvDsInferParseYolo
...
```
**NOTE**: The **RT-DETR Ultralytics** do not resize the input with padding. To get better accuracy, use
```
[property]
...
maintain-aspect-ratio=0
...
```
##
### Edit the deepstream_app_config file
```
...
[primary-gie]
...
config-file=config_infer_primary_rtdetr.txt
```
##
### Testing the model
```
deepstream-app -c deepstream_app_config.txt
```
**NOTE**: The TensorRT engine file may take a very long time to generate (sometimes more than 10 minutes).
**NOTE**: For more information about custom models configuration (`batch-size`, `network-mode`, etc), please check the [`docs/customModels.md`](customModels.md) file.

124
utils/export_rtdetr_ultralytics.py Executable file
View File

@@ -0,0 +1,124 @@
import os
import sys
import argparse
import warnings
import onnx
import torch
import torch.nn as nn
from copy import deepcopy
from ultralytics import RTDETR
from ultralytics.utils.torch_utils import select_device
from ultralytics.nn.modules import C2f, Detect, RTDETRDecoder
class DeepStreamOutput(nn.Module):
def __init__(self, img_size):
self.img_size = img_size
super().__init__()
def forward(self, x):
boxes = x[:, :, :4]
boxes[:, :, [0, 2]] *= self.img_size[1]
boxes[:, :, [1, 3]] *= self.img_size[0]
scores, classes = torch.max(x[:, :, 4:], 2, keepdim=True)
classes = classes.float()
return boxes, scores, classes
def suppress_warnings():
warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
warnings.filterwarnings('ignore', category=UserWarning)
warnings.filterwarnings('ignore', category=DeprecationWarning)
def rtdetr_ultralytics_export(weights, device):
model = RTDETR(weights)
model = deepcopy(model.model).to(device)
for p in model.parameters():
p.requires_grad = False
model.eval()
model.float()
model = model.fuse()
for k, m in model.named_modules():
if isinstance(m, (Detect, RTDETRDecoder)):
m.dynamic = False
m.export = True
m.format = 'onnx'
elif isinstance(m, C2f):
m.forward = m.forward_split
return model
def main(args):
suppress_warnings()
print('\nStarting: %s' % args.weights)
print('Opening RT-DETR Ultralytics model\n')
device = select_device('cpu')
model = rtdetr_ultralytics_export(args.weights, device)
if len(model.names.keys()) > 0:
print('\nCreating labels.txt file')
f = open('labels.txt', 'w')
for name in model.names.values():
f.write(name + '\n')
f.close()
img_size = args.size * 2 if len(args.size) == 1 else args.size
model = nn.Sequential(model, DeepStreamOutput(img_size))
onnx_input_im = torch.zeros(args.batch, 3, *img_size).to(device)
onnx_output_file = os.path.basename(args.weights).split('.pt')[0] + '.onnx'
dynamic_axes = {
'input': {
0: 'batch'
},
'boxes': {
0: 'batch'
},
'scores': {
0: 'batch'
},
'classes': {
0: 'batch'
}
}
print('\nExporting the model to ONNX')
torch.onnx.export(model, onnx_input_im, onnx_output_file, verbose=False, opset_version=args.opset,
do_constant_folding=True, input_names=['input'], output_names=['boxes', 'scores', 'classes'],
dynamic_axes=dynamic_axes if args.dynamic else None)
if args.simplify:
print('Simplifying the ONNX model')
import onnxsim
model_onnx = onnx.load(onnx_output_file)
model_onnx, _ = onnxsim.simplify(model_onnx)
onnx.save(model_onnx, onnx_output_file)
print('Done: %s\n' % onnx_output_file)
def parse_args():
parser = argparse.ArgumentParser(description='DeepStream RT-DETR Ultralytics conversion')
parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
parser.add_argument('-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
parser.add_argument('--opset', type=int, default=16, help='ONNX opset version')
parser.add_argument('--simplify', action='store_true', help='ONNX simplify model')
parser.add_argument('--dynamic', action='store_true', help='Dynamic batch-size')
parser.add_argument('--batch', type=int, default=1, help='Static batch-size')
args = parser.parse_args()
if not os.path.isfile(args.weights):
raise SystemExit('Invalid weights file')
if args.dynamic and args.batch > 1:
raise SystemExit('Cannot set dynamic batch-size and static batch-size at same time')
return args
if __name__ == '__main__':
args = parse_args()
sys.exit(main(args))