New DetectMultiBackend() class (#5549)
* New `DetectMultiBackend()` class * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * pb to pt fix * Cleanup * explicit apply_classifier path * Cleanup2 * Cleanup3 * Cleanup4 * Cleanup5 * Cleanup6 * val.py MultiBackend inference * warmup fix * to device fix * pt fix * device fix * Val cleanup * COCO128 URL to assets * half fix * detect fix * detect fix 2 * remove half from DetectMultiBackend * training half handling * training half handling 2 * training half handling 3 * Cleanup * Fix CI error * Add torchscript _extra_files * Add TorchScript * Add CoreML * CoreML cleanup * New `DetectMultiBackend()` class * pb to pt fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cleanup * explicit apply_classifier path * Cleanup2 * Cleanup3 * Cleanup4 * Cleanup5 * Cleanup6 * val.py MultiBackend inference * warmup fix * to device fix * pt fix * device fix * Val cleanup * COCO128 URL to assets * half fix * detect fix * detect fix 2 * remove half from DetectMultiBackend * training half handling * training half handling 2 * training half handling 3 * Cleanup * Fix CI error * Add torchscript _extra_files * Add TorchScript * Add CoreML * CoreML cleanup * revert default to pt * Add Usage examples * Cleanup val Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocherandpre-commit-ci[bot] authored and GitHub committed
1 parent
79bca2bf64
commit
3883261143
7 files changed
+201
-166
No files matched your search
+127
-1
@@ -3,11 +3,14 @@
|
||||
Common modules
|
||||
"""
|
||||
|
||||
import json
|
||||
import math
|
||||
import platform
|
||||
import warnings
|
||||
from copy import copy
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import requests
|
||||
@@ -17,7 +20,8 @@ from PIL import Image
|
||||
from torch.cuda import amp
|
||||
|
||||
from utils.datasets import exif_transpose, letterbox
|
||||
from utils.general import LOGGER, colorstr, increment_path, make_divisible, non_max_suppression, scale_coords, xyxy2xywh
|
||||
from utils.general import (LOGGER, check_requirements, check_suffix, colorstr, increment_path, make_divisible,
|
||||
non_max_suppression, scale_coords, xywh2xyxy, xyxy2xywh)
|
||||
from utils.plots import Annotator, colors, save_one_box
|
||||
from utils.torch_utils import time_sync
|
||||
|
||||
@@ -269,6 +273,128 @@ class Concat(nn.Module):
|
||||
return torch.cat(x, self.d)
|
||||
|
||||
|
||||
class DetectMultiBackend(nn.Module):
|
||||
# YOLOv5 MultiBackend class for python inference on various backends
|
||||
def __init__(self, weights='yolov5s.pt', device=None, dnn=True):
|
||||
# Usage:
|
||||
# PyTorch: weights = *.pt
|
||||
# TorchScript: *.torchscript.pt
|
||||
# CoreML: *.mlmodel
|
||||
# TensorFlow: *_saved_model
|
||||
# TensorFlow: *.pb
|
||||
# TensorFlow Lite: *.tflite
|
||||
# ONNX Runtime: *.onnx
|
||||
# OpenCV DNN: *.onnx with dnn=True
|
||||
super().__init__()
|
||||
w = str(weights[0] if isinstance(weights, list) else weights)
|
||||
suffix, suffixes = Path(w).suffix.lower(), ['.pt', '.onnx', '.tflite', '.pb', '', '.mlmodel']
|
||||
check_suffix(w, suffixes) # check weights have acceptable suffix
|
||||
pt, onnx, tflite, pb, saved_model, coreml = (suffix == x for x in suffixes) # backend booleans
|
||||
jit = pt and 'torchscript' in w.lower()
|
||||
stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults
|
||||
|
||||
if jit: # TorchScript
|
||||
LOGGER.info(f'Loading {w} for TorchScript inference...')
|
||||
extra_files = {'config.txt': ''} # model metadata
|
||||
model = torch.jit.load(w, _extra_files=extra_files)
|
||||
if extra_files['config.txt']:
|
||||
d = json.loads(extra_files['config.txt']) # extra_files dict
|
||||
stride, names = int(d['stride']), d['names']
|
||||
elif pt: # PyTorch
|
||||
from models.experimental import attempt_load # scoped to avoid circular import
|
||||
model = torch.jit.load(w) if 'torchscript' in w else attempt_load(weights, map_location=device)
|
||||
stride = int(model.stride.max()) # model stride
|
||||
names = model.module.names if hasattr(model, 'module') else model.names # get class names
|
||||
elif coreml: # CoreML *.mlmodel
|
||||
import coremltools as ct
|
||||
model = ct.models.MLModel(w)
|
||||
elif dnn: # ONNX OpenCV DNN
|
||||
LOGGER.info(f'Loading {w} for ONNX OpenCV DNN inference...')
|
||||
check_requirements(('opencv-python>=4.5.4',))
|
||||
net = cv2.dnn.readNetFromONNX(w)
|
||||
elif onnx: # ONNX Runtime
|
||||
LOGGER.info(f'Loading {w} for ONNX Runtime inference...')
|
||||
check_requirements(('onnx', 'onnxruntime-gpu' if torch.has_cuda else 'onnxruntime'))
|
||||
import onnxruntime
|
||||
session = onnxruntime.InferenceSession(w, None)
|
||||
else: # TensorFlow model (TFLite, pb, saved_model)
|
||||
import tensorflow as tf
|
||||
if pb: # https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
|
||||
def wrap_frozen_graph(gd, inputs, outputs):
|
||||
x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), []) # wrapped
|
||||
return x.prune(tf.nest.map_structure(x.graph.as_graph_element, inputs),
|
||||
tf.nest.map_structure(x.graph.as_graph_element, outputs))
|
||||
|
||||
LOGGER.info(f'Loading {w} for TensorFlow *.pb inference...')
|
||||
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...')
|
||||
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():
|
||||
LOGGER.info(f'Loading {w} for TensorFlow 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:
|
||||
LOGGER.info(f'Loading {w} for TensorFlow Lite inference...')
|
||||
interpreter = tf.lite.Interpreter(model_path=w) # load TFLite model
|
||||
interpreter.allocate_tensors() # allocate
|
||||
input_details = interpreter.get_input_details() # inputs
|
||||
output_details = interpreter.get_output_details() # outputs
|
||||
self.__dict__.update(locals()) # assign all variables to self
|
||||
|
||||
def forward(self, im, augment=False, visualize=False, val=False):
|
||||
# YOLOv5 MultiBackend inference
|
||||
b, ch, h, w = im.shape # batch, channel, height, width
|
||||
if self.pt: # 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 *.mlmodel
|
||||
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
|
||||
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]
|
||||
else: # TensorFlow model (TFLite, pb, saved_model)
|
||||
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:
|
||||
y = self.model(im, training=False).numpy()
|
||||
elif self.tflite:
|
||||
input, output = self.input_details[0], self.output_details[0]
|
||||
int8 = input['dtype'] == np.uint8 # is TFLite quantized uint8 model
|
||||
if int8:
|
||||
scale, zero_point = input['quantization']
|
||||
im = (im / scale + zero_point).astype(np.uint8) # de-scale
|
||||
self.interpreter.set_tensor(input['index'], im)
|
||||
self.interpreter.invoke()
|
||||
y = self.interpreter.get_tensor(output['index'])
|
||||
if int8:
|
||||
scale, zero_point = output['quantization']
|
||||
y = (y.astype(np.float32) - zero_point) * scale # re-scale
|
||||
y[..., 0] *= w # x
|
||||
y[..., 1] *= h # y
|
||||
y[..., 2] *= w # w
|
||||
y[..., 3] *= h # h
|
||||
y = torch.tensor(y)
|
||||
return (y, []) if val else y
|
||||
|
||||
|
||||
class AutoShape(nn.Module):
|
||||
# YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
|
||||
conf = 0.25 # NMS confidence threshold
|
||||
|
||||
Reference in new issue
Block a user