├── README.md ├── imgs ├── bus.jpg ├── coco.names ├── dog.jpg ├── person.jpg ├── street.png └── zidane.jpg ├── main.cpp ├── main.py └── onnxmodel ├── v5lite-c.onnx ├── v5lite-g.onnx └── v5lite-s.onnx /README.md: -------------------------------------------------------------------------------- 1 | # yolov5-lite-onnxruntime 2 | 使用ONNXRuntime部署yolov5-lite目标检测，包含C++和Python两个版本的程序 3 | 4 | 起初，我打算使用opencv的部署yolov5-lite，可是opencv的dnn模块读取.onnx文件 5 | 出错了。因此使用ONNXRuntime部署yolov5-lite目标检测 6 | -------------------------------------------------------------------------------- /imgs/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-lite-onnxruntime/0ff8aea2d75fa5cae42ea6fd5fa81ce955808f5a/imgs/bus.jpg -------------------------------------------------------------------------------- /imgs/coco.names: -------------------------------------------------------------------------------- 1 | person 2 | bicycle 3 | car 4 | motorbike 5 | aeroplane 6 | bus 7 | train 8 | truck 9 | boat 10 | traffic light 11 | fire hydrant 12 | stop sign 13 | parking meter 14 | bench 15 | bird 16 | cat 17 | dog 18 | horse 19 | sheep 20 | cow 21 | elephant 22 | bear 23 | zebra 24 | giraffe 25 | backpack 26 | umbrella 27 | handbag 28 | tie 29 | suitcase 30 | frisbee 31 | skis 32 | snowboard 33 | sports ball 34 | kite 35 | baseball bat 36 | baseball glove 37 | skateboard 38 | surfboard 39 | tennis racket 40 | bottle 41 | wine glass 42 | cup 43 | fork 44 | knife 45 | spoon 46 | bowl 47 | banana 48 | apple 49 | sandwich 50 | orange 51 | broccoli 52 | carrot 53 | hot dog 54 | pizza 55 | donut 56 | cake 57 | chair 58 | sofa 59 | pottedplant 60 | bed 61 | diningtable 62 | toilet 63 | tvmonitor 64 | laptop 65 | mouse 66 | remote 67 | keyboard 68 | cell phone 69 | microwave 70 | oven 71 | toaster 72 | sink 73 | refrigerator 74 | book 75 | clock 76 | vase 77 | scissors 78 | teddy bear 79 | hair drier 80 | toothbrush 81 | -------------------------------------------------------------------------------- /imgs/dog.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-lite-onnxruntime/0ff8aea2d75fa5cae42ea6fd5fa81ce955808f5a/imgs/dog.jpg -------------------------------------------------------------------------------- /imgs/person.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-lite-onnxruntime/0ff8aea2d75fa5cae42ea6fd5fa81ce955808f5a/imgs/person.jpg -------------------------------------------------------------------------------- /imgs/street.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-lite-onnxruntime/0ff8aea2d75fa5cae42ea6fd5fa81ce955808f5a/imgs/street.png -------------------------------------------------------------------------------- /imgs/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-lite-onnxruntime/0ff8aea2d75fa5cae42ea6fd5fa81ce955808f5a/imgs/zidane.jpg -------------------------------------------------------------------------------- /main.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-lite-onnxruntime/0ff8aea2d75fa5cae42ea6fd5fa81ce955808f5a/main.cpp -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import argparse 4 | import onnxruntime as ort 5 | import math 6 | 7 | class yolov5_lite(): 8 | def __init__(self, model_pb_path, label_path, confThreshold=0.5, nmsThreshold=0.5, objThreshold=0.5): 9 | so = ort.SessionOptions() 10 | so.log_severity_level = 3 11 | self.net = ort.InferenceSession(model_pb_path, so) 12 | self.classes = list(map(lambda x: x.strip(), open(label_path, 'r').readlines())) 13 | self.num_classes = len(self.classes) 14 | anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]] 15 | self.nl = len(anchors) 16 | self.na = len(anchors[0]) // 2 17 | self.no = self.num_classes + 5 18 | self.grid = [np.zeros(1)] * self.nl 19 | self.stride = np.array([8., 16., 32.]) 20 | self.anchor_grid = np.asarray(anchors, dtype=np.float32).reshape(self.nl, -1, 2) 21 | 22 | self.confThreshold = confThreshold 23 | self.nmsThreshold = nmsThreshold 24 | self.objThreshold = objThreshold 25 | self.input_shape = (self.net.get_inputs()[0].shape[2], self.net.get_inputs()[0].shape[3]) 26 | 27 | def resize_image(self, srcimg, keep_ratio=True): 28 | top, left, newh, neww = 0, 0, self.input_shape[0], self.input_shape[1] 29 | if keep_ratio and srcimg.shape[0] != srcimg.shape[1]: 30 | hw_scale = srcimg.shape[0] / srcimg.shape[1] 31 | if hw_scale > 1: 32 | newh, neww = self.input_shape[0], int(self.input_shape[1] / hw_scale) 33 | img = cv2.resize(srcimg, (neww, newh), interpolation=cv2.INTER_AREA) 34 | left = int((self.input_shape[1] - neww) * 0.5) 35 | img = cv2.copyMakeBorder(img, 0, 0, left, self.input_shape[1] - neww - left, cv2.BORDER_CONSTANT, 36 | value=0) # add border 37 | else: 38 | newh, neww = int(self.input_shape[0] * hw_scale), self.input_shape[1] 39 | img = cv2.resize(srcimg, (neww, newh), interpolation=cv2.INTER_AREA) 40 | top = int((self.input_shape[0] - newh) * 0.5) 41 | img = cv2.copyMakeBorder(img, top, self.input_shape[0] - newh - top, 0, 0, cv2.BORDER_CONSTANT, value=0) 42 | else: 43 | img = cv2.resize(srcimg, self.input_shape, interpolation=cv2.INTER_AREA) 44 | return img, newh, neww, top, left 45 | def _make_grid(self, nx=20, ny=20): 46 | xv, yv = np.meshgrid(np.arange(ny), np.arange(nx)) 47 | return np.stack((xv, yv), 2).reshape((-1, 2)).astype(np.float32) 48 | 49 | def postprocess(self, frame, outs, pad_hw): 50 | newh, neww, padh, padw = pad_hw 51 | frameHeight = frame.shape[0] 52 | frameWidth = frame.shape[1] 53 | ratioh, ratiow = frameHeight / newh, frameWidth / neww 54 | # Scan through all the bounding boxes output from the network and keep only the 55 | # ones with high confidence scores. Assign the box's class label as the class with the highest score. 56 | classIds = [] 57 | confidences = [] 58 | boxes = [] 59 | for detection in outs: 60 | scores = detection[5:] 61 | classId = np.argmax(scores) 62 | confidence = scores[classId] 63 | if confidence > self.confThreshold and detection[4] > self.objThreshold: 64 | center_x = int((detection[0] - padw) * ratiow) 65 | center_y = int((detection[1] - padh) * ratioh) 66 | width = int(detection[2] * ratiow) 67 | height = int(detection[3] * ratioh) 68 | left = int(center_x - width / 2) 69 | top = int(center_y - height / 2) 70 | classIds.append(classId) 71 | confidences.append(float(confidence)) 72 | boxes.append([left, top, width, height]) 73 | 74 | # Perform non maximum suppression to eliminate redundant overlapping boxes with 75 | # lower confidences. 76 | indices = cv2.dnn.NMSBoxes(boxes, confidences, self.confThreshold, self.nmsThreshold) 77 | for i in indices: 78 | i = i[0] if isinstance(i, (tuple,list)) else i 79 | box = boxes[i] 80 | left = box[0] 81 | top = box[1] 82 | width = box[2] 83 | height = box[3] 84 | frame = self.drawPred(frame, classIds[i], confidences[i], left, top, left + width, top + height) 85 | return frame 86 | def drawPred(self, frame, classId, conf, left, top, right, bottom): 87 | # Draw a bounding box. 88 | cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), thickness=4) 89 | 90 | label = '%.2f' % conf 91 | label = '%s:%s' % (self.classes[classId], label) 92 | 93 | # Display the label at the top of the bounding box 94 | labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) 95 | top = max(top, labelSize[1]) 96 | # cv.rectangle(frame, (left, top - round(1.5 * labelSize[1])), (left + round(1.5 * labelSize[0]), top + baseLine), (255,255,255), cv.FILLED) 97 | cv2.putText(frame, label, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), thickness=2) 98 | return frame 99 | def detect(self, srcimg): 100 | img, newh, neww, top, left = self.resize_image(srcimg) 101 | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 102 | img = img.astype(np.float32) / 255.0 103 | blob = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0) 104 | 105 | outs = self.net.run(None, {self.net.get_inputs()[0].name: blob})[0].squeeze(axis=0) 106 | row_ind = 0 107 | for i in range(self.nl): 108 | h, w = int(self.input_shape[0] / self.stride[i]), int(self.input_shape[1] / self.stride[i]) 109 | length = int(self.na * h * w) 110 | if self.grid[i].shape[2:4] != (h, w): 111 | self.grid[i] = self._make_grid(w, h) 112 | 113 | outs[row_ind:row_ind + length, 0:2] = (outs[row_ind:row_ind + length, 0:2] * 2. - 0.5 + np.tile( 114 | self.grid[i], (self.na, 1))) * int(self.stride[i]) 115 | outs[row_ind:row_ind + length, 2:4] = (outs[row_ind:row_ind + length, 2:4] * 2) ** 2 * np.repeat( 116 | self.anchor_grid[i], h * w, axis=0) 117 | row_ind += length 118 | srcimg = self.postprocess(srcimg, outs, (newh, neww, top, left)) 119 | # cv2.imwrite('result.jpg', srcimg) 120 | return srcimg 121 | 122 | if __name__=='__main__': 123 | parser = argparse.ArgumentParser() 124 | parser.add_argument('--imgpath', type=str, default='imgs/horse.jpg', help="image path") 125 | parser.add_argument('--modelpath', type=str, default='onnxmodel/v5lite-s.onnx', help="onnx filepath") 126 | parser.add_argument('--classfile', type=str, default='imgs/coco.names', help="classname filepath") 127 | parser.add_argument('--confThreshold', default=0.5, type=float, help='class confidence') 128 | parser.add_argument('--nmsThreshold', default=0.6, type=float, help='nms iou thresh') 129 | args = parser.parse_args() 130 | 131 | srcimg = cv2.imread(args.imgpath) 132 | net = yolov5_lite(args.modelpath, args.classfile, confThreshold=args.confThreshold, nmsThreshold=args.nmsThreshold) 133 | srcimg = net.detect(srcimg) 134 | 135 | winName = 'Deep learning object detection in onnxruntime' 136 | cv2.namedWindow(winName, cv2.WINDOW_NORMAL) 137 | cv2.imshow(winName, srcimg) 138 | cv2.waitKey(0) 139 | cv2.destroyAllWindows() -------------------------------------------------------------------------------- /onnxmodel/v5lite-c.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-lite-onnxruntime/0ff8aea2d75fa5cae42ea6fd5fa81ce955808f5a/onnxmodel/v5lite-c.onnx -------------------------------------------------------------------------------- /onnxmodel/v5lite-g.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-lite-onnxruntime/0ff8aea2d75fa5cae42ea6fd5fa81ce955808f5a/onnxmodel/v5lite-g.onnx -------------------------------------------------------------------------------- /onnxmodel/v5lite-s.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-lite-onnxruntime/0ff8aea2d75fa5cae42ea6fd5fa81ce955808f5a/onnxmodel/v5lite-s.onnx --------------------------------------------------------------------------------