Add OpenVINO inference (#6179)
Browse files
detect.py
CHANGED
@@ -16,7 +16,7 @@ Usage - formats:
|
|
16 |
yolov5s.torchscript # TorchScript
|
17 |
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
18 |
yolov5s.mlmodel # CoreML (under development)
|
19 |
-
|
20 |
yolov5s_saved_model # TensorFlow SavedModel
|
21 |
yolov5s.pb # TensorFlow protobuf
|
22 |
yolov5s.tflite # TensorFlow Lite
|
|
|
16 |
yolov5s.torchscript # TorchScript
|
17 |
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
18 |
yolov5s.mlmodel # CoreML (under development)
|
19 |
+
yolov5s.xml # OpenVINO
|
20 |
yolov5s_saved_model # TensorFlow SavedModel
|
21 |
yolov5s.pb # TensorFlow protobuf
|
22 |
yolov5s.tflite # TensorFlow Lite
|
export.py
CHANGED
@@ -20,16 +20,16 @@ Usage:
|
|
20 |
$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml openvino saved_model tflite tfjs
|
21 |
|
22 |
Inference:
|
23 |
-
$ python path/to/detect.py --weights yolov5s.pt
|
24 |
-
yolov5s.torchscript
|
25 |
-
yolov5s.onnx
|
26 |
-
yolov5s.mlmodel
|
27 |
-
|
28 |
-
yolov5s_saved_model
|
29 |
-
yolov5s.pb
|
30 |
-
yolov5s.tflite
|
31 |
-
yolov5s_edgetpu.tflite
|
32 |
-
yolov5s.engine
|
33 |
|
34 |
TensorFlow.js:
|
35 |
$ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
|
|
|
20 |
$ python path/to/export.py --weights yolov5s.pt --include torchscript onnx coreml openvino saved_model tflite tfjs
|
21 |
|
22 |
Inference:
|
23 |
+
$ python path/to/detect.py --weights yolov5s.pt # PyTorch
|
24 |
+
yolov5s.torchscript # TorchScript
|
25 |
+
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
26 |
+
yolov5s.mlmodel # CoreML (under development)
|
27 |
+
yolov5s.xml # OpenVINO
|
28 |
+
yolov5s_saved_model # TensorFlow SavedModel
|
29 |
+
yolov5s.pb # TensorFlow protobuf
|
30 |
+
yolov5s.tflite # TensorFlow Lite
|
31 |
+
yolov5s_edgetpu.tflite # TensorFlow Edge TPU
|
32 |
+
yolov5s.engine # TensorRT
|
33 |
|
34 |
TensorFlow.js:
|
35 |
$ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
|
models/common.py
CHANGED
@@ -282,6 +282,7 @@ class DetectMultiBackend(nn.Module):
|
|
282 |
# PyTorch: weights = *.pt
|
283 |
# TorchScript: *.torchscript
|
284 |
# CoreML: *.mlmodel
|
|
|
285 |
# TensorFlow: *_saved_model
|
286 |
# TensorFlow: *.pb
|
287 |
# TensorFlow Lite: *.tflite
|
@@ -294,31 +295,38 @@ class DetectMultiBackend(nn.Module):
|
|
294 |
super().__init__()
|
295 |
w = str(weights[0] if isinstance(weights, list) else weights)
|
296 |
suffix = Path(w).suffix.lower()
|
297 |
-
suffixes = ['.pt', '.torchscript', '.onnx', '.engine', '.tflite', '.pb', '', '.mlmodel']
|
298 |
check_suffix(w, suffixes) # check weights have acceptable suffix
|
299 |
-
pt, jit, onnx, engine, tflite, pb, saved_model, coreml = (suffix == x for x in suffixes) #
|
300 |
stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
|
301 |
w = attempt_download(w) # download if not local
|
302 |
if data: # data.yaml path (optional)
|
303 |
with open(data, errors='ignore') as f:
|
304 |
names = yaml.safe_load(f)['names'] # class names
|
305 |
|
306 |
-
if
|
|
|
|
|
|
|
|
|
|
|
307 |
LOGGER.info(f'Loading {w} for TorchScript inference...')
|
308 |
extra_files = {'config.txt': ''} # model metadata
|
309 |
model = torch.jit.load(w, _extra_files=extra_files)
|
310 |
if extra_files['config.txt']:
|
311 |
d = json.loads(extra_files['config.txt']) # extra_files dict
|
312 |
stride, names = int(d['stride']), d['names']
|
313 |
-
elif pt: # PyTorch
|
314 |
-
model = attempt_load(weights if isinstance(weights, list) else w, map_location=device)
|
315 |
-
stride = int(model.stride.max()) # model stride
|
316 |
-
names = model.module.names if hasattr(model, 'module') else model.names # get class names
|
317 |
-
self.model = model # explicitly assign for to(), cpu(), cuda(), half()
|
318 |
elif coreml: # CoreML
|
319 |
LOGGER.info(f'Loading {w} for CoreML inference...')
|
320 |
import coremltools as ct
|
321 |
model = ct.models.MLModel(w)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
322 |
elif dnn: # ONNX OpenCV DNN
|
323 |
LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...')
|
324 |
check_requirements(('opencv-python>=4.5.4',))
|
@@ -403,6 +411,13 @@ class DetectMultiBackend(nn.Module):
|
|
403 |
y = self.net.forward()
|
404 |
else: # ONNX Runtime
|
405 |
y = self.session.run([self.session.get_outputs()[0].name], {self.session.get_inputs()[0].name: im})[0]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
406 |
elif self.engine: # TensorRT
|
407 |
assert im.shape == self.bindings['images'].shape, (im.shape, self.bindings['images'].shape)
|
408 |
self.binding_addrs['images'] = int(im.data_ptr())
|
|
|
282 |
# PyTorch: weights = *.pt
|
283 |
# TorchScript: *.torchscript
|
284 |
# CoreML: *.mlmodel
|
285 |
+
# OpenVINO: *.xml
|
286 |
# TensorFlow: *_saved_model
|
287 |
# TensorFlow: *.pb
|
288 |
# TensorFlow Lite: *.tflite
|
|
|
295 |
super().__init__()
|
296 |
w = str(weights[0] if isinstance(weights, list) else weights)
|
297 |
suffix = Path(w).suffix.lower()
|
298 |
+
suffixes = ['.pt', '.torchscript', '.onnx', '.engine', '.tflite', '.pb', '', '.mlmodel', '.xml']
|
299 |
check_suffix(w, suffixes) # check weights have acceptable suffix
|
300 |
+
pt, jit, onnx, engine, tflite, pb, saved_model, coreml, xml = (suffix == x for x in suffixes) # backends
|
301 |
stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
|
302 |
w = attempt_download(w) # download if not local
|
303 |
if data: # data.yaml path (optional)
|
304 |
with open(data, errors='ignore') as f:
|
305 |
names = yaml.safe_load(f)['names'] # class names
|
306 |
|
307 |
+
if pt: # PyTorch
|
308 |
+
model = attempt_load(weights if isinstance(weights, list) else w, map_location=device)
|
309 |
+
stride = int(model.stride.max()) # model stride
|
310 |
+
names = model.module.names if hasattr(model, 'module') else model.names # get class names
|
311 |
+
self.model = model # explicitly assign for to(), cpu(), cuda(), half()
|
312 |
+
elif jit: # TorchScript
|
313 |
LOGGER.info(f'Loading {w} for TorchScript inference...')
|
314 |
extra_files = {'config.txt': ''} # model metadata
|
315 |
model = torch.jit.load(w, _extra_files=extra_files)
|
316 |
if extra_files['config.txt']:
|
317 |
d = json.loads(extra_files['config.txt']) # extra_files dict
|
318 |
stride, names = int(d['stride']), d['names']
|
|
|
|
|
|
|
|
|
|
|
319 |
elif coreml: # CoreML
|
320 |
LOGGER.info(f'Loading {w} for CoreML inference...')
|
321 |
import coremltools as ct
|
322 |
model = ct.models.MLModel(w)
|
323 |
+
elif xml: # OpenVINO
|
324 |
+
LOGGER.info(f'Loading {w} for OpenVINO inference...')
|
325 |
+
check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/
|
326 |
+
import openvino.inference_engine as ie
|
327 |
+
core = ie.IECore()
|
328 |
+
network = core.read_network(model=w, weights=Path(w).with_suffix('.bin')) # *.xml, *.bin paths
|
329 |
+
executable_network = core.load_network(network, device_name='CPU', num_requests=1)
|
330 |
elif dnn: # ONNX OpenCV DNN
|
331 |
LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...')
|
332 |
check_requirements(('opencv-python>=4.5.4',))
|
|
|
411 |
y = self.net.forward()
|
412 |
else: # ONNX Runtime
|
413 |
y = self.session.run([self.session.get_outputs()[0].name], {self.session.get_inputs()[0].name: im})[0]
|
414 |
+
elif self.xml: # OpenVINO
|
415 |
+
im = im.cpu().numpy() # FP32
|
416 |
+
desc = self.ie.TensorDesc(precision='FP32', dims=im.shape, layout='NCHW') # Tensor Description
|
417 |
+
request = self.executable_network.requests[0] # inference request
|
418 |
+
request.set_blob(blob_name='images', blob=self.ie.Blob(desc, im)) # name=next(iter(request.input_blobs))
|
419 |
+
request.infer()
|
420 |
+
y = request.output_blobs['output'].buffer # name=next(iter(request.output_blobs))
|
421 |
elif self.engine: # TensorRT
|
422 |
assert im.shape == self.bindings['images'].shape, (im.shape, self.bindings['images'].shape)
|
423 |
self.binding_addrs['images'] = int(im.data_ptr())
|
val.py
CHANGED
@@ -10,7 +10,7 @@ Usage - formats:
|
|
10 |
yolov5s.torchscript # TorchScript
|
11 |
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
12 |
yolov5s.mlmodel # CoreML (under development)
|
13 |
-
|
14 |
yolov5s_saved_model # TensorFlow SavedModel
|
15 |
yolov5s.pb # TensorFlow protobuf
|
16 |
yolov5s.tflite # TensorFlow Lite
|
|
|
10 |
yolov5s.torchscript # TorchScript
|
11 |
yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn
|
12 |
yolov5s.mlmodel # CoreML (under development)
|
13 |
+
yolov5s.xml # OpenVINO
|
14 |
yolov5s_saved_model # TensorFlow SavedModel
|
15 |
yolov5s.pb # TensorFlow protobuf
|
16 |
yolov5s.tflite # TensorFlow Lite
|