├── README.md ├── main.cpp ├── main.py ├── s_l.jpg ├── selfie.jpg └── test.jpg /README.md: -------------------------------------------------------------------------------- 1 | # yolov7-detect-face-onnxrun-cpp-py 2 | 使用ONNXRuntime部署YOLOV7人脸+关键点检测，包含C++和Python两个版本的程序 3 | 4 | 起初想使用OpenCV做部署的，但是OpenCV的dnn模块读取onnx文件出错了，无奈只能使用onnxruntime做部署 5 | 6 | 7 | onnx模型文件在百度云盘，链接：https://pan.baidu.com/s/1oIj3QgOohKpHGwJDxNvDsw 8 | 提取码：ebv8 9 | 10 | yolov7人脸+关键点检测的训练源码，在 11 | https://github.com/derronqi/yolov7-face 12 | -------------------------------------------------------------------------------- /main.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov7-detect-face-onnxrun-cpp-py/7aad2f4c9c9e0f607ce7335df2800d68e284d0e2/main.cpp -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import onnxruntime 4 | import argparse 5 | 6 | 7 | class YOLOv7_face: 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 = ['face'] 12 | # Initialize model 13 | session_option = onnxruntime.SessionOptions() 14 | session_option.log_severity_level = 3 15 | # self.session = onnxruntime.InferenceSession(path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) 16 | self.session = onnxruntime.InferenceSession(path, sess_options=session_option) 17 | model_inputs = self.session.get_inputs() 18 | self.input_names = [model_inputs[i].name for i in range(len(model_inputs))] 19 | self.input_shape = model_inputs[0].shape 20 | self.input_height = int(self.input_shape[2]) 21 | self.input_width = int(self.input_shape[3]) 22 | 23 | model_outputs = self.session.get_outputs() 24 | self.output_names = [model_outputs[i].name for i in range(len(model_outputs))] 25 | 26 | def resize_image(self, srcimg, keep_ratio=True): 27 | top, left, newh, neww = 0, 0, self.input_width, self.input_height 28 | if keep_ratio and srcimg.shape[0] != srcimg.shape[1]: 29 | hw_scale = srcimg.shape[0] / srcimg.shape[1] 30 | if hw_scale > 1: 31 | newh, neww = self.input_height, int(self.input_width / hw_scale) 32 | img = cv2.resize(srcimg, (neww, newh), interpolation=cv2.INTER_AREA) 33 | left = int((self.input_width - neww) * 0.5) 34 | img = cv2.copyMakeBorder(img, 0, 0, left, self.input_width - neww - left, cv2.BORDER_CONSTANT, 35 | value=(114, 114, 114)) # add border 36 | else: 37 | newh, neww = int(self.input_height * hw_scale), self.input_width 38 | img = cv2.resize(srcimg, (neww, newh), interpolation=cv2.INTER_AREA) 39 | top = int((self.input_height - newh) * 0.5) 40 | img = cv2.copyMakeBorder(img, top, self.input_height - newh - top, 0, 0, cv2.BORDER_CONSTANT, 41 | value=(114, 114, 114)) 42 | else: 43 | img = cv2.resize(srcimg, (self.input_width, self.input_height), interpolation=cv2.INTER_AREA) 44 | return img, newh, neww, top, left 45 | 46 | def prepare_input(self, image): 47 | self.img_height, self.img_width = image.shape[:2] 48 | self.scale = np.array( 49 | [self.img_width / self.input_width, self.img_height / self.input_height, self.img_width / self.input_width, 50 | self.img_height / self.input_height], dtype=np.float32) 51 | input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 52 | 53 | # Resize input image 54 | input_img = cv2.resize(input_img, (self.input_width, self.input_height)) 55 | # input_img, newh, neww, top, left = self.resize_image(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) ###也可以使用保持高宽比resize的pad填充 56 | 57 | # Scale input pixel values to 0 to 1 58 | input_img = input_img.astype(np.float32) / 255.0 59 | input_img = input_img.transpose(2, 0, 1) 60 | input_tensor = input_img[np.newaxis, :, :, :] 61 | return input_tensor 62 | 63 | def detect(self, image): 64 | input_tensor = self.prepare_input(image) 65 | 66 | # Perform inference on the image 67 | outputs = self.session.run(self.output_names, {input_name: input_tensor for input_name in self.input_names}) 68 | boxes, scores, kpts = self.process_output(outputs) 69 | return boxes, scores, kpts 70 | 71 | def process_output(self, output): 72 | predictions = np.squeeze(output[0]).reshape((-1, 21)) 73 | 74 | # Filter out object confidence scores below threshold 75 | obj_conf = predictions[:, 4] 76 | predictions = predictions[obj_conf > self.conf_threshold] 77 | obj_conf = obj_conf[obj_conf > self.conf_threshold] 78 | 79 | # Multiply class confidence with bounding box confidence 80 | predictions[:, 5] *= obj_conf 81 | 82 | # Get the scores 83 | scores = predictions[:, 5] 84 | 85 | # Filter out the objects with a low score 86 | valid_scores = scores > self.conf_threshold 87 | predictions = predictions[valid_scores] 88 | scores = scores[valid_scores] 89 | 90 | # Get bounding boxes for each object 91 | boxes, kpts = self.extract_boxes(predictions) 92 | indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), self.conf_threshold, 93 | self.iou_threshold) 94 | return boxes[indices], scores[indices], kpts[indices] 95 | 96 | def extract_boxes(self, predictions): 97 | # Extract boxes from predictions 98 | boxes = predictions[:, :4] * self.scale 99 | kpts = predictions[:, 6:] ###x1,y1,score1, ...., x5,y5,score5 100 | kpts *= np.tile(np.array([self.scale[0], self.scale[1], 1], dtype=np.float32), (1, 5)) 101 | 102 | # Convert boxes to xywh format 103 | boxes_ = np.copy(boxes) 104 | boxes_[..., 0] = boxes[..., 0] - boxes[..., 2] * 0.5 105 | boxes_[..., 1] = boxes[..., 1] - boxes[..., 3] * 0.5 106 | return boxes_, kpts 107 | 108 | def draw_detections(self, image, boxes, scores, kpts): 109 | for box, score, kp in zip(boxes, scores, kpts): 110 | x, y, w, h = box.astype(int) 111 | 112 | # Draw rectangle 113 | cv2.rectangle(image, (x, y), (x + w, y + h), (0, 0, 255), thickness=2) 114 | label = self.class_names[0] 115 | labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) 116 | # top = max(y1, labelSize[1]) 117 | # cv.rectangle(frame, (left, top - round(1.5 * labelSize[1])), (left + round(1.5 * labelSize[0]), top + baseLine), (255,255,255), cv.FILLED) 118 | cv2.putText(image, label, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), thickness=2) 119 | for i in range(5): 120 | cv2.circle(image, (int(kp[i * 3]), int(kp[i * 3 + 1])), 1, (0, 255, 0), thickness=-1) 121 | return image 122 | 123 | 124 | if __name__ == '__main__': 125 | parser = argparse.ArgumentParser() 126 | parser.add_argument('--imgpath', type=str, default='selfie.jpg', help="image path") 127 | parser.add_argument('--modelpath', type=str, default='onnx_havepost_models/yolov7-lite-e.onnx', 128 | help="onnx filepath") 129 | parser.add_argument('--confThreshold', default=0.45, type=float, help='class confidence') 130 | parser.add_argument('--nmsThreshold', default=0.5, type=float, help='nms iou thresh') 131 | args = parser.parse_args() 132 | 133 | # Initialize YOLOv7_face object detector 134 | YOLOv7_face_detector = YOLOv7_face(args.modelpath, conf_thres=args.confThreshold, iou_thres=args.nmsThreshold) 135 | srcimg = cv2.imread(args.imgpath) 136 | 137 | # Detect Objects 138 | boxes, scores, kpts = YOLOv7_face_detector.detect(srcimg) 139 | 140 | # Draw detections 141 | dstimg = YOLOv7_face_detector.draw_detections(srcimg, boxes, scores, kpts) 142 | winName = 'Deep learning object detection in ONNXRuntime' 143 | cv2.namedWindow(winName, 0) 144 | cv2.imshow(winName, dstimg) 145 | cv2.waitKey(0) 146 | cv2.destroyAllWindows() 147 | -------------------------------------------------------------------------------- /s_l.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov7-detect-face-onnxrun-cpp-py/7aad2f4c9c9e0f607ce7335df2800d68e284d0e2/s_l.jpg -------------------------------------------------------------------------------- /selfie.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov7-detect-face-onnxrun-cpp-py/7aad2f4c9c9e0f607ce7335df2800d68e284d0e2/selfie.jpg -------------------------------------------------------------------------------- /test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hpc203/yolov7-detect-face-onnxrun-cpp-py/7aad2f4c9c9e0f607ce7335df2800d68e284d0e2/test.jpg --------------------------------------------------------------------------------