Glenn Jocher authored and GitHub committed
1 parent
7cad6597bb
commit
5bd6a97b18
4 files changed
+120
-120
No files matched your search
+40
-40
@@ -316,17 +316,6 @@ class DetectMultiBackend(nn.Module):
|
||||
if extra_files['config.txt']:
|
||||
d = json.loads(extra_files['config.txt']) # extra_files dict
|
||||
stride, names = int(d['stride']), d['names']
|
||||
elif coreml: # CoreML
|
||||
LOGGER.info(f'Loading {w} for CoreML inference...')
|
||||
import coremltools as ct
|
||||
model = ct.models.MLModel(w)
|
||||
elif xml: # OpenVINO
|
||||
LOGGER.info(f'Loading {w} for OpenVINO inference...')
|
||||
check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/
|
||||
import openvino.inference_engine as ie
|
||||
core = ie.IECore()
|
||||
network = core.read_network(model=w, weights=Path(w).with_suffix('.bin')) # *.xml, *.bin paths
|
||||
executable_network = core.load_network(network, device_name='CPU', num_requests=1)
|
||||
elif dnn: # ONNX OpenCV DNN
|
||||
LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...')
|
||||
check_requirements(('opencv-python>=4.5.4',))
|
||||
@@ -338,6 +327,13 @@ class DetectMultiBackend(nn.Module):
|
||||
import onnxruntime
|
||||
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
|
||||
session = onnxruntime.InferenceSession(w, providers=providers)
|
||||
elif xml: # OpenVINO
|
||||
LOGGER.info(f'Loading {w} for OpenVINO inference...')
|
||||
check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/
|
||||
import openvino.inference_engine as ie
|
||||
core = ie.IECore()
|
||||
network = core.read_network(model=w, weights=Path(w).with_suffix('.bin')) # *.xml, *.bin paths
|
||||
executable_network = core.load_network(network, device_name='CPU', num_requests=1)
|
||||
elif engine: # TensorRT
|
||||
LOGGER.info(f'Loading {w} for TensorRT inference...')
|
||||
import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download
|
||||
@@ -356,9 +352,17 @@ class DetectMultiBackend(nn.Module):
|
||||
binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
|
||||
context = model.create_execution_context()
|
||||
batch_size = bindings['images'].shape[0]
|
||||
else: # TensorFlow (TFLite, pb, saved_model)
|
||||
if pb: # https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
|
||||
LOGGER.info(f'Loading {w} for TensorFlow *.pb inference...')
|
||||
elif coreml: # CoreML
|
||||
LOGGER.info(f'Loading {w} for CoreML inference...')
|
||||
import coremltools as ct
|
||||
model = ct.models.MLModel(w)
|
||||
else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
|
||||
if saved_model: # SavedModel
|
||||
LOGGER.info(f'Loading {w} for TensorFlow SavedModel inference...')
|
||||
import tensorflow as tf
|
||||
model = tf.keras.models.load_model(w)
|
||||
elif pb: # GraphDef https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
|
||||
LOGGER.info(f'Loading {w} for TensorFlow GraphDef inference...')
|
||||
import tensorflow as tf
|
||||
|
||||
def wrap_frozen_graph(gd, inputs, outputs):
|
||||
@@ -369,19 +373,15 @@ class DetectMultiBackend(nn.Module):
|
||||
graph_def = tf.Graph().as_graph_def()
|
||||
graph_def.ParseFromString(open(w, 'rb').read())
|
||||
frozen_func = wrap_frozen_graph(gd=graph_def, inputs="x:0", outputs="Identity:0")
|
||||
elif saved_model:
|
||||
LOGGER.info(f'Loading {w} for TensorFlow saved_model inference...')
|
||||
import tensorflow as tf
|
||||
model = tf.keras.models.load_model(w)
|
||||
elif tflite: # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python
|
||||
if 'edgetpu' in w.lower():
|
||||
if 'edgetpu' in w.lower(): # Edge TPU
|
||||
LOGGER.info(f'Loading {w} for TensorFlow Lite Edge TPU inference...')
|
||||
import tflite_runtime.interpreter as tfli
|
||||
delegate = {'Linux': 'libedgetpu.so.1', # install https://coral.ai/software/#edgetpu-runtime
|
||||
'Darwin': 'libedgetpu.1.dylib',
|
||||
'Windows': 'edgetpu.dll'}[platform.system()]
|
||||
interpreter = tfli.Interpreter(model_path=w, experimental_delegates=[tfli.load_delegate(delegate)])
|
||||
else:
|
||||
else: # Lite
|
||||
LOGGER.info(f'Loading {w} for TensorFlow Lite inference...')
|
||||
import tensorflow as tf
|
||||
interpreter = tf.lite.Interpreter(model_path=w) # load TFLite model
|
||||
@@ -396,21 +396,13 @@ class DetectMultiBackend(nn.Module):
|
||||
if self.pt or self.jit: # PyTorch
|
||||
y = self.model(im) if self.jit else self.model(im, augment=augment, visualize=visualize)
|
||||
return y if val else y[0]
|
||||
elif self.coreml: # CoreML
|
||||
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
|
||||
im = Image.fromarray((im[0] * 255).astype('uint8'))
|
||||
# im = im.resize((192, 320), Image.ANTIALIAS)
|
||||
y = self.model.predict({'image': im}) # coordinates are xywh normalized
|
||||
box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels
|
||||
conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)
|
||||
y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
|
||||
elif self.onnx: # ONNX
|
||||
elif self.dnn: # ONNX OpenCV DNN
|
||||
im = im.cpu().numpy() # torch to numpy
|
||||
if self.dnn: # ONNX OpenCV DNN
|
||||
self.net.setInput(im)
|
||||
y = self.net.forward()
|
||||
else: # ONNX Runtime
|
||||
y = self.session.run([self.session.get_outputs()[0].name], {self.session.get_inputs()[0].name: im})[0]
|
||||
self.net.setInput(im)
|
||||
y = self.net.forward()
|
||||
elif self.onnx: # ONNX Runtime
|
||||
im = im.cpu().numpy() # torch to numpy
|
||||
y = self.session.run([self.session.get_outputs()[0].name], {self.session.get_inputs()[0].name: im})[0]
|
||||
elif self.xml: # OpenVINO
|
||||
im = im.cpu().numpy() # FP32
|
||||
desc = self.ie.TensorDesc(precision='FP32', dims=im.shape, layout='NCHW') # Tensor Description
|
||||
@@ -423,13 +415,21 @@ class DetectMultiBackend(nn.Module):
|
||||
self.binding_addrs['images'] = int(im.data_ptr())
|
||||
self.context.execute_v2(list(self.binding_addrs.values()))
|
||||
y = self.bindings['output'].data
|
||||
else: # TensorFlow model (TFLite, pb, saved_model)
|
||||
elif self.coreml: # CoreML
|
||||
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
|
||||
if self.pb:
|
||||
y = self.frozen_func(x=self.tf.constant(im)).numpy()
|
||||
elif self.saved_model:
|
||||
im = Image.fromarray((im[0] * 255).astype('uint8'))
|
||||
# im = im.resize((192, 320), Image.ANTIALIAS)
|
||||
y = self.model.predict({'image': im}) # coordinates are xywh normalized
|
||||
box = xywh2xyxy(y['coordinates'] * [[w, h, w, h]]) # xyxy pixels
|
||||
conf, cls = y['confidence'].max(1), y['confidence'].argmax(1).astype(np.float)
|
||||
y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
|
||||
else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
|
||||
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
|
||||
if self.saved_model: # SavedModel
|
||||
y = self.model(im, training=False).numpy()
|
||||
elif self.tflite:
|
||||
elif self.pb: # GraphDef
|
||||
y = self.frozen_func(x=self.tf.constant(im)).numpy()
|
||||
elif self.tflite: # Lite
|
||||
input, output = self.input_details[0], self.output_details[0]
|
||||
int8 = input['dtype'] == np.uint8 # is TFLite quantized uint8 model
|
||||
if int8:
|
||||
@@ -451,7 +451,7 @@ class DetectMultiBackend(nn.Module):
|
||||
|
||||
def warmup(self, imgsz=(1, 3, 640, 640), half=False):
|
||||
# Warmup model by running inference once
|
||||
if self.pt or self.engine or self.onnx: # warmup types
|
||||
if self.pt or self.jit or self.onnx or self.engine: # warmup types
|
||||
if isinstance(self.device, torch.device) and self.device.type != 'cpu': # only warmup GPU models
|
||||
im = torch.zeros(*imgsz).to(self.device).type(torch.half if half else torch.float) # input image
|
||||
self.forward(im) # warmup
|
||||
|
||||
Reference in new issue
Block a user