Add OpenCV DNN option for ONNX inference (#5136)
Browse files* Add OpenCV DNN option for ONNX inference
Usage:
```bash
python detect.py --weights yolov5s.onnx # ONNX Runtime inference
python detect.py --weights yolov5s.onnx -dnn # OpenCV DNN inference
```
* DNN prediction to tensor
* Update detect.py
detect.py
CHANGED
@@ -56,6 +56,7 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
|
|
56 |
hide_labels=False, # hide labels
|
57 |
hide_conf=False, # hide confidences
|
58 |
half=False, # use FP16 half-precision inference
|
|
|
59 |
):
|
60 |
source = str(source)
|
61 |
save_img = not nosave and not source.endswith('.txt') # save inference images
|
@@ -72,7 +73,7 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
|
|
72 |
half &= device.type != 'cpu' # half precision only supported on CUDA
|
73 |
|
74 |
# Load model
|
75 |
-
w = weights[0] if isinstance(weights, list) else weights
|
76 |
classify, suffix, suffixes = False, Path(w).suffix.lower(), ['.pt', '.onnx', '.tflite', '.pb', '']
|
77 |
check_suffix(w, suffixes) # check weights have acceptable suffix
|
78 |
pt, onnx, tflite, pb, saved_model = (suffix == x for x in suffixes) # backend booleans
|
@@ -87,9 +88,13 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
|
|
87 |
modelc = load_classifier(name='resnet50', n=2) # initialize
|
88 |
modelc.load_state_dict(torch.load('resnet50.pt', map_location=device)['model']).to(device).eval()
|
89 |
elif onnx:
|
90 |
-
|
91 |
-
|
92 |
-
|
|
|
|
|
|
|
|
|
93 |
else: # TensorFlow models
|
94 |
check_requirements(('tensorflow>=2.4.1',))
|
95 |
import tensorflow as tf
|
@@ -145,7 +150,11 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
|
|
145 |
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
|
146 |
pred = model(img, augment=augment, visualize=visualize)[0]
|
147 |
elif onnx:
|
148 |
-
|
|
|
|
|
|
|
|
|
149 |
else: # tensorflow model (tflite, pb, saved_model)
|
150 |
imn = img.permute(0, 2, 3, 1).cpu().numpy() # image in numpy
|
151 |
if pb:
|
@@ -281,6 +290,7 @@ def parse_opt():
|
|
281 |
parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
|
282 |
parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
|
283 |
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
|
|
|
284 |
opt = parser.parse_args()
|
285 |
opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
|
286 |
print_args(FILE.stem, opt)
|
|
|
56 |
hide_labels=False, # hide labels
|
57 |
hide_conf=False, # hide confidences
|
58 |
half=False, # use FP16 half-precision inference
|
59 |
+
dnn=False, # use OpenCV DNN for ONNX inference
|
60 |
):
|
61 |
source = str(source)
|
62 |
save_img = not nosave and not source.endswith('.txt') # save inference images
|
|
|
73 |
half &= device.type != 'cpu' # half precision only supported on CUDA
|
74 |
|
75 |
# Load model
|
76 |
+
w = str(weights[0] if isinstance(weights, list) else weights)
|
77 |
classify, suffix, suffixes = False, Path(w).suffix.lower(), ['.pt', '.onnx', '.tflite', '.pb', '']
|
78 |
check_suffix(w, suffixes) # check weights have acceptable suffix
|
79 |
pt, onnx, tflite, pb, saved_model = (suffix == x for x in suffixes) # backend booleans
|
|
|
88 |
modelc = load_classifier(name='resnet50', n=2) # initialize
|
89 |
modelc.load_state_dict(torch.load('resnet50.pt', map_location=device)['model']).to(device).eval()
|
90 |
elif onnx:
|
91 |
+
if dnn:
|
92 |
+
# check_requirements(('opencv-python>=4.5.4',))
|
93 |
+
net = cv2.dnn.readNetFromONNX(w)
|
94 |
+
else:
|
95 |
+
check_requirements(('onnx', 'onnxruntime'))
|
96 |
+
import onnxruntime
|
97 |
+
session = onnxruntime.InferenceSession(w, None)
|
98 |
else: # TensorFlow models
|
99 |
check_requirements(('tensorflow>=2.4.1',))
|
100 |
import tensorflow as tf
|
|
|
150 |
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
|
151 |
pred = model(img, augment=augment, visualize=visualize)[0]
|
152 |
elif onnx:
|
153 |
+
if dnn:
|
154 |
+
net.setInput(img)
|
155 |
+
pred = torch.tensor(net.forward())
|
156 |
+
else:
|
157 |
+
pred = torch.tensor(session.run([session.get_outputs()[0].name], {session.get_inputs()[0].name: img}))
|
158 |
else: # tensorflow model (tflite, pb, saved_model)
|
159 |
imn = img.permute(0, 2, 3, 1).cpu().numpy() # image in numpy
|
160 |
if pb:
|
|
|
290 |
parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
|
291 |
parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
|
292 |
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
|
293 |
+
parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
|
294 |
opt = parser.parse_args()
|
295 |
opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
|
296 |
print_args(FILE.stem, opt)
|