YOLOV5 onnxruntime预测不用pytorch
onnxruntime
microsoft/onnxruntime: 是一个用于运行各种机器学习模型的开源库。适合对机器学习和深度学习有兴趣的人,特别是在开发和部署机器学习模型时需要处理各种不同框架和算子的人。特点是支持多种机器学习框架和算子,包括 TensorFlow、PyTorch、Caffe 等,具有高性能和广泛的兼容性。
项目地址:https://gitcode.com/gh_mirrors/on/onnxruntime
免费下载资源
·
import numpy as np
import cv2
import onnxruntime
import argparse
def letterbox(img, new_shape=(416, 416), color=(114, 114, 114), auto=False, scaleFill=False, scaleup=True):
shape = img.shape[:2] # current shape [height, width]
if isinstance(new_shape, int):
new_shape = (new_shape, new_shape)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
if not scaleup:
r = min(r, 1.0)
ratio = r, r # width, height ratios
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
if auto: # minimum rectangle
dw, dh = np.mod(dw, 64), np.mod(dh, 64) # wh padding
elif scaleFill: # stretch
dw, dh = 0.0, 0.0
new_unpad = (new_shape[1], new_shape[0])
ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios
dw /= 2 # divide padding into 2 sides
dh /= 2
if shape[::-1] != new_unpad: # resize
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
return img, ratio, (dw, dh)
def clip_coords(boxes, img_shape):
# Clip bounding xyxy bounding boxes to image shape (height, width)
boxes[:, 0].clip(0, img_shape[1]) # x1
boxes[:, 1].clip(0, img_shape[0]) # y1
boxes[:, 2].clip(0, img_shape[1]) # x2
boxes[:, 3].clip(0, img_shape[0]) # y2
def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
# Rescale coords (xyxy) from img1_shape to img0_shape
if ratio_pad is None: # calculate from img0_shape
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
else:
gain = ratio_pad[0][0]
pad = ratio_pad[1]
coords[:, [0, 2]] -= pad[0] # x padding
coords[:, [1, 3]] -= pad[1] # y padding
coords[:, :4] /= gain
clip_coords(coords, img0_shape)
return coords
class Detector():
def __init__(self, opt):
super(Detector, self).__init__()
self.img_size = opt.img_size
self.threshold = opt.conf_thres
self.iou_thres = opt.iou_thres
self.stride = 1
self.weights = opt.weights
self.init_model()
self.names = ["wu","yzm"]
self.detected_labels = ["wu","yzm"]
def init_model(self):
sess = onnxruntime.InferenceSession(self.weights)
self.input_name = sess.get_inputs()[0].name
output_names = []
for i in range(len(sess.get_outputs())):
print('output shape:', sess.get_outputs()[i].name)
output_names.append(sess.get_outputs()[i].name)
self.output_name = sess.get_outputs()[0].name
print('input name:%s, output name:%s' % (self.input_name, self.output_name))
input_shape = sess.get_inputs()[0].shape
print('input_shape:', input_shape)
self.m = sess
def preprocess(self, img):
img0 = img.copy()
img = letterbox(img, new_shape=self.img_size)[0]
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img).astype(np.float32)
img /= 255.0 # 图像归一化
img = np.expand_dims(img, axis=0)
assert len(img.shape) == 4
return img0, img
def detect(self, im):
im0, img = self.preprocess(im)
W, H = img.shape[2:]
pred = self.m.run(None, {self.input_name: img})[0]
pred = pred.astype(np.float32)
pred = np.squeeze(pred, axis=0)
boxes = []
classIds = []
confidences = []
for detection in pred:
scores = detection[5:]
classID = np.argmax(scores)
confidence = scores[classID] * detection[4]
if confidence > self.threshold:
box = detection[0:4]
(centerX, centerY, width, height) = box.astype("int")
x = int(centerX - (width / 2))
y = int(centerY - (height / 2))
boxes.append([x, y, int(width), int(height)])
classIds.append(classID)
confidences.append(float(confidence))
idxs = cv2.dnn.NMSBoxes(
boxes, confidences, self.threshold, self.iou_thres)
pred_boxes = []
pred_confes = []
pred_classes = []
if len(idxs) > 0:
for i in idxs.flatten():
confidence = confidences[i]
if confidence >= self.threshold:
pred_boxes.append(boxes[i])
pred_confes.append(confidence)
pred_classes.append(classIds[i])
return im, pred_boxes, pred_confes, pred_classes
def main(opt):
det = Detector(opt)
image = cv2.imread(opt.source)
shape = (det.img_size, det.img_size)
im0, pred_boxes, pred_confes, pred_classes = det.detect(image)
if len(pred_boxes) > 0:
for i, _ in enumerate(pred_boxes):
box = pred_boxes[i]
left, top, width, height = box[0], box[1], box[2], box[3]
box = (left, top, left + width, top + height)
box = np.squeeze(
scale_coords(shape, np.expand_dims(box, axis=0).astype("float"), im0.shape[:2]).round(),
axis=0).astype("int")
x0, y0, x1, y1 = box[0], box[1], box[2], box[3]
cv2.rectangle(image, (x0, y0), (x1, y1), (0, 0, 255), thickness=2)
cv2.putText(image, '{0}--{1:.2f}'.format(pred_classes[i], pred_confes[i]), (x0, y0 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), thickness=1)
cv2.imshow("detector", image)
cv2.waitKey()
cv2.destroyAllWindows()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--weights', nargs='+', type=str, default='detect_model.onnx', help='onnx path(s)')
parser.add_argument('--source', type=str, default='yzm.jpg', help='source') # file/folder, 0 for webcam
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
parser.add_argument('--line-thickness', default=1, type=int, help='bounding box thickness (pixels)')
parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
opt = parser.parse_args()
#print(opt)
main(opt)
这个版本完全使用onnxruntime,这样可以减少环境的大小。
GitHub 加速计划 / on / onnxruntime
13.76 K
2.79 K
下载
microsoft/onnxruntime: 是一个用于运行各种机器学习模型的开源库。适合对机器学习和深度学习有兴趣的人,特别是在开发和部署机器学习模型时需要处理各种不同框架和算子的人。特点是支持多种机器学习框架和算子,包括 TensorFlow、PyTorch、Caffe 等,具有高性能和广泛的兼容性。
最近提交(Master分支:26 天前 )
bd60add8
### Description
<!-- Describe your changes. -->
Use the latest nuget.exe for the `readme` property to be supported.
### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
#22137 8 小时前
99ee6eec
### Description
<!-- Describe your changes. -->
Add linker flags to support 16KB page size support on Android.
See
https://source.android.com/docs/core/architecture/16kb-page-size/16kb#build-lib-16kb-alignment
### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
#21837 8 小时前
更多推荐
已为社区贡献1条内容
所有评论(0)