├── README.md ├── images ├── gather_137918.jpg ├── gather_138383.jpg ├── gather_159086.jpg ├── gather_159847.jpg └── gather_87804.jpg ├── main.cpp └── main.py /README.md: -------------------------------------------------------------------------------- 1 | # yolov7-head-onnxrun-cpp-py 2 | 使用ONNXRuntime部署YOLOV7人头检测，包含C++和Python两个版本的程序. 3 | 起初想使用opencv部署的，可是opencv读取onnx文件后在forward函数出错了， 4 | 无赖只能使用onnxruntime部署。 5 | 这套程序对onnxruntime的版本要求较高，需要使用最新的onnxruntime做推理。 6 | 我的机器上的onnxruntime版本是1.11.1，人头检测可以应用到人流密度估计场景里。 7 | 8 | onnx文件在百度云盘，链接：https://pan.baidu.com/s/1N22OMJA1UVMpao2QT9Y4lw 9 | 提取码：t02a 10 | -------------------------------------------------------------------------------- /images/gather_137918.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov7-head-detect-onnxrun-cpp-py/034bfc1ace49ade584132677680f2be0ecf0cf01/images/gather_137918.jpg -------------------------------------------------------------------------------- /images/gather_138383.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov7-head-detect-onnxrun-cpp-py/034bfc1ace49ade584132677680f2be0ecf0cf01/images/gather_138383.jpg -------------------------------------------------------------------------------- /images/gather_159086.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov7-head-detect-onnxrun-cpp-py/034bfc1ace49ade584132677680f2be0ecf0cf01/images/gather_159086.jpg -------------------------------------------------------------------------------- /images/gather_159847.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov7-head-detect-onnxrun-cpp-py/034bfc1ace49ade584132677680f2be0ecf0cf01/images/gather_159847.jpg -------------------------------------------------------------------------------- /images/gather_87804.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov7-head-detect-onnxrun-cpp-py/034bfc1ace49ade584132677680f2be0ecf0cf01/images/gather_87804.jpg -------------------------------------------------------------------------------- /main.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov7-head-detect-onnxrun-cpp-py/034bfc1ace49ade584132677680f2be0ecf0cf01/main.cpp -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import onnxruntime 4 | import argparse 5 | 6 | 7 | class YOLOv7: 8 | def __init__(self, path, conf_thres=0.2, iou_thres=0.5): 9 | self.conf_threshold = conf_thres 10 | self.iou_threshold = iou_thres 11 | self.class_names = ['head'] 12 | # Initialize model 13 | session_option = onnxruntime.SessionOptions() 14 | session_option.log_severity_level = 3 15 | self.session = onnxruntime.InferenceSession(path, sess_options=session_option) 16 | model_inputs = self.session.get_inputs() 17 | self.input_names = [model_inputs[i].name for i in range(len(model_inputs))] 18 | self.input_shape = model_inputs[0].shape 19 | self.input_height = int(self.input_shape[2]) 20 | self.input_width = int(self.input_shape[3]) 21 | 22 | model_outputs = self.session.get_outputs() 23 | self.output_names = [model_outputs[i].name for i in range(len(model_outputs))] 24 | self.has_postprocess = False if len(self.output_names)==1 else True 25 | print(self.has_postprocess) 26 | 27 | def prepare_input(self, image): 28 | self.img_height, self.img_width = image.shape[:2] 29 | 30 | input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 31 | 32 | # Resize input image 33 | input_img = cv2.resize(input_img, (self.input_width, self.input_height)) 34 | 35 | # Scale input pixel values to 0 to 1 36 | input_img = input_img.astype(np.float32) / 255.0 37 | input_img = input_img.transpose(2, 0, 1) 38 | input_tensor = input_img[np.newaxis, :, :, :].astype(np.float32) 39 | return input_tensor 40 | 41 | def detect(self, image): 42 | input_tensor = self.prepare_input(image) 43 | 44 | # Perform inference on the image 45 | outputs = self.session.run(self.output_names, {input_name: input_tensor for input_name in self.input_names}) 46 | 47 | if self.has_postprocess: 48 | boxes, scores, class_ids = self.parse_processed_output(outputs) 49 | 50 | else: 51 | # Process output data 52 | boxes, scores, class_ids = self.process_output(outputs) 53 | 54 | return boxes, scores, class_ids 55 | 56 | 57 | def process_output(self, output): 58 | predictions = np.squeeze(output[0]) 59 | 60 | # Filter out object confidence scores below threshold 61 | obj_conf = predictions[:, 4] 62 | predictions = predictions[obj_conf > self.conf_threshold] 63 | obj_conf = obj_conf[obj_conf > self.conf_threshold] 64 | 65 | # Multiply class confidence with bounding box confidence 66 | predictions[:, 5:] *= obj_conf[:, np.newaxis] 67 | 68 | # Get the scores 69 | scores = np.max(predictions[:, 5:], axis=1) 70 | 71 | # Filter out the objects with a low score 72 | valid_scores = scores > self.conf_threshold 73 | predictions = predictions[valid_scores] 74 | scores = scores[valid_scores] 75 | 76 | # Get the class with the highest confidence 77 | class_ids = np.argmax(predictions[:, 5:], axis=1) 78 | 79 | # Get bounding boxes for each object 80 | boxes = self.extract_boxes(predictions) 81 | indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), self.conf_threshold, 82 | self.iou_threshold).flatten() 83 | return boxes[indices], scores[indices], class_ids[indices] 84 | 85 | def parse_processed_output(self, outputs): 86 | scores = np.squeeze(outputs[0]) 87 | predictions = outputs[1] 88 | 89 | # Filter out object scores below threshold 90 | valid_scores = scores > self.conf_threshold 91 | predictions = predictions[valid_scores, :] 92 | scores = scores[valid_scores] 93 | 94 | # Extract the boxes and class ids 95 | class_ids = predictions[:, 1] 96 | boxes = predictions[:, 2:] 97 | 98 | # In postprocess, the x,y are the y,x 99 | boxes = boxes[:, [1, 0, 3, 2]] 100 | boxes = self.rescale_boxes(boxes) 101 | return boxes, scores, class_ids 102 | 103 | def extract_boxes(self, predictions): 104 | # Extract boxes from predictions 105 | boxes = predictions[:, :4] 106 | 107 | # Scale boxes to original image dimensions 108 | boxes = self.rescale_boxes(boxes) 109 | 110 | # Convert boxes to xywh format 111 | boxes_ = np.copy(boxes) 112 | boxes_[..., 0] = boxes[..., 0] - boxes[..., 2] * 0.5 113 | boxes_[..., 1] = boxes[..., 1] - boxes[..., 3] * 0.5 114 | return boxes_ 115 | 116 | def rescale_boxes(self, boxes): 117 | 118 | # Rescale boxes to original image dimensions 119 | input_shape = np.array([self.input_width, self.input_height, self.input_width, self.input_height]) 120 | boxes = np.divide(boxes, input_shape, dtype=np.float32) 121 | boxes *= np.array([self.img_width, self.img_height, self.img_width, self.img_height]) 122 | return boxes 123 | 124 | def draw_detections(self, image, boxes, scores, class_ids): 125 | for box, score, class_id in zip(boxes, scores, class_ids): 126 | x, y, w, h = box.astype(int) 127 | 128 | # Draw rectangle 129 | cv2.rectangle(image, (x, y), (x+w, y+h), (0, 0, 255), thickness=2) 130 | label = self.class_names[class_id] 131 | label = f'{label} {int(score * 100)}%' 132 | labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) 133 | # top = max(y1, labelSize[1]) 134 | # cv.rectangle(frame, (left, top - round(1.5 * labelSize[1])), (left + round(1.5 * labelSize[0]), top + baseLine), (255,255,255), cv.FILLED) 135 | cv2.putText(image, label, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), thickness=2) 136 | return image 137 | 138 | 139 | if __name__ == '__main__': 140 | parser = argparse.ArgumentParser() 141 | parser.add_argument('--imgpath', type=str, default='images/gather_159847.jpg', help="image path") 142 | parser.add_argument('--modelpath', type=str, default='models/yolov7_head_0.752_480x640.onnx', help="onnx filepath") 143 | parser.add_argument('--confThreshold', default=0.2, type=float, help='class confidence') 144 | parser.add_argument('--nmsThreshold', default=0.5, type=float, help='nms iou thresh') 145 | args = parser.parse_args() 146 | 147 | # Initialize YOLOv7 object detector 148 | yolov7_detector = YOLOv7(args.modelpath, conf_thres=args.confThreshold, iou_thres=args.nmsThreshold) 149 | srcimg = cv2.imread(args.imgpath) 150 | 151 | # Detect Objects 152 | boxes, scores, class_ids = yolov7_detector.detect(srcimg) 153 | 154 | # Draw detections 155 | dstimg = yolov7_detector.draw_detections(srcimg, boxes, scores, class_ids) 156 | winName = 'Deep learning object detection in ONNXRuntime' 157 | cv2.namedWindow(winName, 0) 158 | cv2.imshow(winName, dstimg) 159 | cv2.waitKey(0) 160 | cv2.destroyAllWindows() 161 | --------------------------------------------------------------------------------