glenn-jocher commited on
Commit
0bf24cf
·
unverified ·
1 Parent(s): 9d75e42

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

Files changed (1) hide show
  1. detect.py +15 -5
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
- check_requirements(('onnx', 'onnxruntime'))
91
- import onnxruntime
92
- session = onnxruntime.InferenceSession(w, None)
 
 
 
 
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
- pred = torch.tensor(session.run([session.get_outputs()[0].name], {session.get_inputs()[0].name: img}))
 
 
 
 
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)