├── README.md ├── onnxruntime ├── class.names ├── images │ ├── bus.jpg │ └── zidane.jpg ├── main.cpp └── main.py └── opencv ├── class.names ├── images ├── bus.jpg └── zidane.jpg ├── main.cpp └── main.py /README.md: -------------------------------------------------------------------------------- 1 | # yolov5-v6.1-opencv-onnxrun 2 | 看到https://github.com/ultralytics/yolov5 在最近更新的v6.1版本的，我编写了一套 3 | 4 | 使用OpenCV部署yolov5-v6.1目标检测，包含C++和Python两个版本的程序 5 | 6 | 使用ONNXRuntime部署yolov5-v6.1目标检测，包含C++和Python两个版本的程序 7 | 8 | 支持yolov5s，yolov5m，yolov5l，yolov5n，yolov5x， 9 | yolov5s6，yolov5m6，yolov5l6，yolov5n6，yolov5x6的十种结构的yolov5-v6.1 10 | 11 | 由于模型文件太多，我没有直接上传到仓库里，模型文件存放在百度云盘里 12 | 13 | 下载链接：https://pan.baidu.com/s/1pz1A4YOurxL5G0_ITVEerg 14 | 提取码：zm2q 15 | 16 | 17 | 转换生成onnx文件的方法，可以参考我的csdn博客文章 18 | https://blog.csdn.net/nihate/article/details/112731327 19 | 或者知乎文章 https://zhuanlan.zhihu.com/p/466677699 20 | -------------------------------------------------------------------------------- /onnxruntime/class.names: -------------------------------------------------------------------------------- 1 | person 2 | bicycle 3 | car 4 | motorcycle 5 | airplane 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 | couch 59 | potted plant 60 | bed 61 | dining table 62 | toilet 63 | tv 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 | -------------------------------------------------------------------------------- /onnxruntime/images/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-v6.1-opencv-onnxrun/5b00418492ffc291c485170bdc3dec6fe172c8a3/onnxruntime/images/bus.jpg -------------------------------------------------------------------------------- /onnxruntime/images/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-v6.1-opencv-onnxrun/5b00418492ffc291c485170bdc3dec6fe172c8a3/onnxruntime/images/zidane.jpg -------------------------------------------------------------------------------- /onnxruntime/main.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-v6.1-opencv-onnxrun/5b00418492ffc291c485170bdc3dec6fe172c8a3/onnxruntime/main.cpp -------------------------------------------------------------------------------- /onnxruntime/main.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import argparse 3 | import numpy as np 4 | import onnxruntime as ort 5 | 6 | class yolov5(): 7 | def __init__(self, modelpath, confThreshold=0.5, nmsThreshold=0.5, objThreshold=0.5): 8 | with open('class.names', 'rt') as f: 9 | self.classes = f.read().rstrip('\n').split('\n') 10 | self.num_classes = len(self.classes) 11 | if modelpath.endswith('6.onnx'): 12 | self.inpHeight, self.inpWidth = 1280, 1280 13 | anchors = [[19, 27, 44, 40, 38, 94], [96, 68, 86, 152, 180, 137], [140, 301, 303, 264, 238, 542], [436, 615, 739, 380, 925, 792]] 14 | self.stride = np.array([8., 16., 32., 64.]) 15 | else: 16 | self.inpHeight, self.inpWidth = 640, 640 17 | anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]] 18 | self.stride = np.array([8., 16., 32.]) 19 | self.nl = len(anchors) 20 | self.na = len(anchors[0]) // 2 21 | self.grid = [np.zeros(1)] * self.nl 22 | self.anchor_grid = np.asarray(anchors, dtype=np.float32).reshape(self.nl, -1, 2) 23 | so = ort.SessionOptions() 24 | so.log_severity_level = 3 25 | self.net = ort.InferenceSession(modelpath, so) 26 | self.confThreshold = confThreshold 27 | self.nmsThreshold = nmsThreshold 28 | self.objThreshold = objThreshold 29 | # self.inpHeight, self.inpWidth = (self.net.get_inputs()[0].shape[2], self.net.get_inputs()[0].shape[3]) 30 | 31 | def resize_image(self, srcimg, keep_ratio=True): 32 | top, left, newh, neww = 0, 0, self.inpWidth, self.inpHeight 33 | if keep_ratio and srcimg.shape[0] != srcimg.shape[1]: 34 | hw_scale = srcimg.shape[0] / srcimg.shape[1] 35 | if hw_scale > 1: 36 | newh, neww = self.inpHeight, int(self.inpWidth / hw_scale) 37 | img = cv2.resize(srcimg, (neww, newh), interpolation=cv2.INTER_AREA) 38 | left = int((self.inpWidth - neww) * 0.5) 39 | img = cv2.copyMakeBorder(img, 0, 0, left, self.inpWidth - neww - left, cv2.BORDER_CONSTANT, 40 | value=(114, 114, 114)) # add border 41 | else: 42 | newh, neww = int(self.inpHeight * hw_scale), self.inpWidth 43 | img = cv2.resize(srcimg, (neww, newh), interpolation=cv2.INTER_AREA) 44 | top = int((self.inpHeight - newh) * 0.5) 45 | img = cv2.copyMakeBorder(img, top, self.inpHeight - newh - top, 0, 0, cv2.BORDER_CONSTANT, 46 | value=(114, 114, 114)) 47 | else: 48 | img = cv2.resize(srcimg, (self.inpWidth, self.inpHeight), interpolation=cv2.INTER_AREA) 49 | return img, newh, neww, top, left 50 | 51 | def _make_grid(self, nx=20, ny=20): 52 | xv, yv = np.meshgrid(np.arange(ny), np.arange(nx)) 53 | return np.stack((xv, yv), 2).reshape((-1, 2)).astype(np.float32) 54 | 55 | def preprocess(self, img): 56 | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 57 | img = img.astype(np.float32) / 255.0 58 | return img 59 | 60 | def postprocess(self, frame, outs, padsize=None): 61 | frameHeight = frame.shape[0] 62 | frameWidth = frame.shape[1] 63 | newh, neww, padh, padw = padsize 64 | ratioh, ratiow = frameHeight / newh, frameWidth / neww 65 | # Scan through all the bounding boxes output from the network and keep only the 66 | # ones with high confidence scores. Assign the box's class label as the class with the highest score. 67 | 68 | confidences = [] 69 | boxes = [] 70 | classIds = [] 71 | for detection in outs: 72 | if detection[4] > self.objThreshold: 73 | scores = detection[5:] 74 | classId = np.argmax(scores) 75 | confidence = scores[classId] * detection[4] 76 | if confidence > self.confThreshold: 77 | center_x = int((detection[0] - padw) * ratiow) 78 | center_y = int((detection[1] - padh) * ratioh) 79 | width = int(detection[2] * ratiow) 80 | height = int(detection[3] * ratioh) 81 | left = int(center_x - width * 0.5) 82 | top = int(center_y - height * 0.5) 83 | 84 | confidences.append(float(confidence)) 85 | boxes.append([left, top, width, height]) 86 | classIds.append(classId) 87 | # Perform non maximum suppression to eliminate redundant overlapping boxes with 88 | # lower confidences. 89 | indices = cv2.dnn.NMSBoxes(boxes, confidences, self.confThreshold, self.nmsThreshold).flatten() 90 | for i in indices: 91 | box = boxes[i] 92 | left = box[0] 93 | top = box[1] 94 | width = box[2] 95 | height = box[3] 96 | frame = self.drawPred(frame, classIds[i], confidences[i], left, top, left + width, top + height) 97 | return frame 98 | 99 | def drawPred(self, frame, classId, conf, left, top, right, bottom): 100 | # Draw a bounding box. 101 | cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), thickness=4) 102 | 103 | label = '%.2f' % conf 104 | label = '%s:%s' % (self.classes[classId], label) 105 | 106 | # Display the label at the top of the bounding box 107 | labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) 108 | top = max(top, labelSize[1]) 109 | # cv.rectangle(frame, (left, top - round(1.5 * labelSize[1])), (left + round(1.5 * labelSize[0]), top + baseLine), (255,255,255), cv.FILLED) 110 | cv2.putText(frame, label, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), thickness=2) 111 | return frame 112 | 113 | def detect(self, srcimg): 114 | img, newh, neww, padh, padw = self.resize_image(srcimg) 115 | img = self.preprocess(img) 116 | # Sets the input to the network 117 | blob = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0) 118 | 119 | outs = self.net.run(None, {self.net.get_inputs()[0].name: blob})[0].squeeze(axis=0) 120 | 121 | # inference output 122 | row_ind = 0 123 | for i in range(self.nl): 124 | h, w = int(img.shape[0] / self.stride[i]), int(img.shape[1] / self.stride[i]) 125 | length = int(self.na * h * w) 126 | if self.grid[i].shape[2:4] != (h, w): 127 | self.grid[i] = self._make_grid(w, h) 128 | 129 | outs[row_ind:row_ind + length, 0:2] = (outs[row_ind:row_ind + length, 0:2] * 2. - 0.5 + np.tile( 130 | self.grid[i], (self.na, 1))) * int(self.stride[i]) 131 | outs[row_ind:row_ind + length, 2:4] = (outs[row_ind:row_ind + length, 2:4] * 2) ** 2 * np.repeat( 132 | self.anchor_grid[i], h * w, axis=0) 133 | row_ind += length 134 | srcimg = self.postprocess(srcimg, outs, padsize=(newh, neww, padh, padw)) 135 | return srcimg 136 | 137 | if __name__ == "__main__": 138 | parser = argparse.ArgumentParser() 139 | parser.add_argument('--imgpath', type=str, default='images/bus.jpg', help="image path") 140 | parser.add_argument('--modelpath', type=str, default='weights/yolov5s.onnx') 141 | parser.add_argument('--confThreshold', default=0.3, type=float, help='class confidence') 142 | parser.add_argument('--nmsThreshold', default=0.5, type=float, help='nms iou thresh') 143 | parser.add_argument('--objThreshold', default=0.3, type=float, help='object confidence') 144 | args = parser.parse_args() 145 | 146 | yolonet = yolov5(args.modelpath, confThreshold=args.confThreshold, nmsThreshold=args.nmsThreshold, 147 | objThreshold=args.objThreshold) 148 | srcimg = cv2.imread(args.imgpath) 149 | srcimg = yolonet.detect(srcimg) 150 | 151 | winName = 'Deep learning object detection in ONNXRuntime' 152 | cv2.namedWindow(winName, 0) 153 | cv2.imshow(winName, srcimg) 154 | cv2.waitKey(0) 155 | cv2.destroyAllWindows() -------------------------------------------------------------------------------- /opencv/class.names: -------------------------------------------------------------------------------- 1 | person 2 | bicycle 3 | car 4 | motorcycle 5 | airplane 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 | couch 59 | potted plant 60 | bed 61 | dining table 62 | toilet 63 | tv 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 | -------------------------------------------------------------------------------- /opencv/images/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-v6.1-opencv-onnxrun/5b00418492ffc291c485170bdc3dec6fe172c8a3/opencv/images/bus.jpg -------------------------------------------------------------------------------- /opencv/images/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-v6.1-opencv-onnxrun/5b00418492ffc291c485170bdc3dec6fe172c8a3/opencv/images/zidane.jpg -------------------------------------------------------------------------------- /opencv/main.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov5-v6.1-opencv-onnxrun/5b00418492ffc291c485170bdc3dec6fe172c8a3/opencv/main.cpp -------------------------------------------------------------------------------- /opencv/main.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import argparse 3 | import numpy as np 4 | 5 | class yolov5(): 6 | def __init__(self, modelpath, confThreshold=0.5, nmsThreshold=0.5, objThreshold=0.5): 7 | with open('class.names', 'rt') as f: 8 | self.classes = f.read().rstrip('\n').split('\n') 9 | self.num_classes = len(self.classes) 10 | if modelpath.endswith('6.onnx'): 11 | self.inpHeight, self.inpWidth = 1280, 1280 12 | anchors = [[19, 27, 44, 40, 38, 94], [96, 68, 86, 152, 180, 137], [140, 301, 303, 264, 238, 542], 13 | [436, 615, 739, 380, 925, 792]] 14 | self.stride = np.array([8., 16., 32., 64.]) 15 | else: 16 | self.inpHeight, self.inpWidth = 640, 640 17 | anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]] 18 | self.stride = np.array([8., 16., 32.]) 19 | self.nl = len(anchors) 20 | self.na = len(anchors[0]) // 2 21 | self.grid = [np.zeros(1)] * self.nl 22 | self.anchor_grid = np.asarray(anchors, dtype=np.float32).reshape(self.nl, -1, 2) 23 | self.net = cv2.dnn.readNet(modelpath) 24 | self.confThreshold = confThreshold 25 | self.nmsThreshold = nmsThreshold 26 | self.objThreshold = objThreshold 27 | self._inputNames = '' 28 | 29 | def resize_image(self, srcimg, keep_ratio=True, dynamic=False): 30 | top, left, newh, neww = 0, 0, self.inpWidth, self.inpHeight 31 | if keep_ratio and srcimg.shape[0] != srcimg.shape[1]: 32 | hw_scale = srcimg.shape[0] / srcimg.shape[1] 33 | if hw_scale > 1: 34 | newh, neww = self.inpHeight, int(self.inpWidth / hw_scale) 35 | img = cv2.resize(srcimg, (neww, newh), interpolation=cv2.INTER_AREA) 36 | if not dynamic: 37 | left = int((self.inpWidth - neww) * 0.5) 38 | img = cv2.copyMakeBorder(img, 0, 0, left, self.inpWidth - neww - left, cv2.BORDER_CONSTANT, 39 | value=(114, 114, 114)) # add border 40 | else: 41 | newh, neww = int(self.inpHeight * hw_scale), self.inpWidth 42 | img = cv2.resize(srcimg, (neww, newh), interpolation=cv2.INTER_AREA) 43 | if not dynamic: 44 | top = int((self.inpHeight - newh) * 0.5) 45 | img = cv2.copyMakeBorder(img, top, self.inpHeight - newh - top, 0, 0, cv2.BORDER_CONSTANT, 46 | value=(114, 114, 114)) 47 | else: 48 | img = cv2.resize(srcimg, (self.inpWidth, self.inpHeight), interpolation=cv2.INTER_AREA) 49 | return img, newh, neww, top, left 50 | 51 | def _make_grid(self, nx=20, ny=20): 52 | xv, yv = np.meshgrid(np.arange(ny), np.arange(nx)) 53 | return np.stack((xv, yv), 2).reshape((-1, 2)).astype(np.float32) 54 | 55 | def preprocess(self, img): 56 | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 57 | img = img.astype(np.float32) / 255.0 58 | return img 59 | 60 | def postprocess(self, frame, outs, padsize=None): 61 | frameHeight = frame.shape[0] 62 | frameWidth = frame.shape[1] 63 | newh, neww, padh, padw = padsize 64 | ratioh, ratiow = frameHeight / newh, frameWidth / neww 65 | # Scan through all the bounding boxes output from the network and keep only the 66 | # ones with high confidence scores. Assign the box's class label as the class with the highest score. 67 | 68 | confidences = [] 69 | boxes = [] 70 | classIds = [] 71 | for detection in outs: 72 | if detection[4] > self.objThreshold: 73 | scores = detection[5:] 74 | classId = np.argmax(scores) 75 | confidence = scores[classId] * detection[4] 76 | if confidence > self.confThreshold: 77 | center_x = int((detection[0] - padw) * ratiow) 78 | center_y = int((detection[1] - padh) * ratioh) 79 | width = int(detection[2] * ratiow) 80 | height = int(detection[3] * ratioh) 81 | left = int(center_x - width * 0.5) 82 | top = int(center_y - height * 0.5) 83 | 84 | confidences.append(float(confidence)) 85 | boxes.append([left, top, width, height]) 86 | classIds.append(classId) 87 | # Perform non maximum suppression to eliminate redundant overlapping boxes with 88 | # lower confidences. 89 | indices = cv2.dnn.NMSBoxes(boxes, confidences, self.confThreshold, self.nmsThreshold).flatten() 90 | for i in indices: 91 | box = boxes[i] 92 | left = box[0] 93 | top = box[1] 94 | width = box[2] 95 | height = box[3] 96 | frame = self.drawPred(frame, classIds[i], confidences[i], left, top, left + width, top + height) 97 | return frame 98 | 99 | def drawPred(self, frame, classId, conf, left, top, right, bottom): 100 | # Draw a bounding box. 101 | cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), thickness=4) 102 | 103 | label = '%.2f' % conf 104 | label = '%s:%s' % (self.classes[classId], label) 105 | 106 | # Display the label at the top of the bounding box 107 | labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) 108 | top = max(top, labelSize[1]) 109 | # cv.rectangle(frame, (left, top - round(1.5 * labelSize[1])), (left + round(1.5 * labelSize[0]), top + baseLine), (255,255,255), cv.FILLED) 110 | cv2.putText(frame, label, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), thickness=2) 111 | return frame 112 | 113 | def detect(self, srcimg): 114 | img, newh, neww, padh, padw = self.resize_image(srcimg) 115 | blob = cv2.dnn.blobFromImage(img, scalefactor=1 / 255.0, swapRB=True) 116 | # blob = cv2.dnn.blobFromImage(self.preprocess(img)) 117 | # Sets the input to the network 118 | self.net.setInput(blob, self._inputNames) 119 | 120 | # Runs the forward pass to get output of the output layers 121 | outs = self.net.forward(self.net.getUnconnectedOutLayersNames())[0].squeeze(axis=0) 122 | 123 | # inference output 124 | row_ind = 0 125 | for i in range(self.nl): 126 | h, w = int(self.inpHeight / self.stride[i]), int(self.inpWidth / self.stride[i]) 127 | length = int(self.na * h * w) 128 | if self.grid[i].shape[2:4] != (h, w): 129 | self.grid[i] = self._make_grid(w, h) 130 | 131 | outs[row_ind:row_ind + length, 0:2] = (outs[row_ind:row_ind + length, 0:2] * 2. - 0.5 + np.tile( 132 | self.grid[i], (self.na, 1))) * int(self.stride[i]) 133 | outs[row_ind:row_ind + length, 2:4] = (outs[row_ind:row_ind + length, 2:4] * 2) ** 2 * np.repeat( 134 | self.anchor_grid[i], h * w, axis=0) 135 | row_ind += length 136 | srcimg = self.postprocess(srcimg, outs, padsize=(newh, neww, padh, padw)) 137 | return srcimg 138 | 139 | if __name__ == "__main__": 140 | parser = argparse.ArgumentParser() 141 | parser.add_argument('--imgpath', type=str, default='images/bus.jpg', help="image path") 142 | parser.add_argument('--modelpath', type=str, default='weights/yolov5s.onnx') 143 | parser.add_argument('--confThreshold', default=0.3, type=float, help='class confidence') 144 | parser.add_argument('--nmsThreshold', default=0.5, type=float, help='nms iou thresh') 145 | parser.add_argument('--objThreshold', default=0.3, type=float, help='object confidence') 146 | args = parser.parse_args() 147 | 148 | yolonet = yolov5(args.modelpath, confThreshold=args.confThreshold, nmsThreshold=args.nmsThreshold, 149 | objThreshold=args.objThreshold) 150 | srcimg = cv2.imread(args.imgpath) 151 | srcimg = yolonet.detect(srcimg) 152 | 153 | winName = 'Deep learning object detection in OpenCV' 154 | cv2.namedWindow(winName, 0) 155 | cv2.imshow(winName, srcimg) 156 | cv2.waitKey(0) 157 | cv2.destroyAllWindows() --------------------------------------------------------------------------------