读取模型

net= cv2.dnn.readNetFromONNX(model_path)

从文件中加载模型

参数说明

model_path：模型文件路径

返回值

对应模型网络的对象

数据预处理

cv2.dnn.blobFromImage(img, scalefactor, size, mean)

输入数据预处理

参数说明

img：待处理图像
scalefactor：用来和图像相乘的比例因子
size：缩放尺寸
mean：用来和图像相减的均值

返回值

用于模型输入的blob对象

设置输入数据

net.setInput(blob)

准备模型网络的输入数据

参数说明

blob：输入blob对象

返回值

无

执行预测

net.forward(output_list)

执行模型推理，并根据输入数组返回对应的网络层

参数说明

output_list：需要返回的网络层

返回值

指定需要返回的网络层

使用示例：打开摄像头，检测画面中的人脸

  1. # coding=utf-8
  2. import os
  3. import time
  4. from math import ceil
  5. import cv2
  6. import numpy as np
  8. Key_Esc = 27
  10. image_mean = np.array([127, 127, 127])
  11. image_std = 128.0
  12. iou_threshold = 0.3
  13. center_variance = 0.1
  14. size_variance = 0.2
  15. min_boxes = [[10.0, 16.0, 24.0], [32.0, 48.0],
  16.              [64.0, 96.0], [128.0, 192.0, 256.0]]
  17. strides = [8.0, 16.0, 32.0, 64.0]
  20. def define_img_size(image_size):
  21.     shrinkage_list = []
  22.     feature_map_w_h_list = []
  23.     for size in image_size:
  24.         feature_map = [int(ceil(size / stride)) for stride in strides]
  25.         feature_map_w_h_list.append(feature_map)
  27.     for i in range(0, len(image_size)):
  28.         shrinkage_list.append(strides)
  29.     priors = generate_priors(feature_map_w_h_list,
  30.                              shrinkage_list, image_size, min_boxes)
  31.     return priors
  34. def generate_priors(feature_map_list, shrinkage_list, image_size, min_boxes):
  35.     priors = []
  36.     for index in range(0, len(feature_map_list[0])):
  37.         scale_w = image_size[0] / shrinkage_list[0][index]
  38.         scale_h = image_size[1] / shrinkage_list[1][index]
  39.         for j in range(0, feature_map_list[1][index]):
  40.             for i in range(0, feature_map_list[0][index]):
  41.                 x_center = (i + 0.5) / scale_w
  42.                 y_center = (j + 0.5) / scale_h
  44.                 for min_box in min_boxes[index]:
  45.                     w = min_box / image_size[0]
  46.                     h = min_box / image_size[1]
  47.                     priors.append([
  48.                         x_center,
  49.                         y_center,
  50.                         w,
  51.                         h
  52.                     ])
  53.     print("priors nums:{}".format(len(priors)))
  54.     return np.clip(priors, 0.0, 1.0)
  57. def hard_nms(box_scores, iou_threshold, top_k=-1, candidate_size=200):
  58.     scores = box_scores[:, -1]
  59.     boxes = box_scores[:, :-1]
  60.     picked = []
  61.     indexes = np.argsort(scores)
  62.     indexes = indexes[-candidate_size:]
  63.     while len(indexes) > 0:
  64.         current = indexes[-1]
  65.         picked.append(current)
  66.         if 0 < top_k == len(picked) or len(indexes) == 1:
  67.             break
  68.         current_box = boxes[current, :]
  69.         indexes = indexes[:-1]
  70.         rest_boxes = boxes[indexes, :]
  71.         iou = iou_of(
  72.             rest_boxes,
  73.             np.expand_dims(current_box, axis=0),
  74.         )
  75.         indexes = indexes[iou <= iou_threshold]
  76.     return box_scores[picked, :]
  79. def area_of(left_top, right_bottom):
  80.     hw = np.clip(right_bottom - left_top, 0.0, None)
  81.     return hw[..., 0] * hw[..., 1]
  84. def iou_of(boxes0, boxes1, eps=1e-5):
  85.     overlap_left_top = np.maximum(boxes0[..., :2], boxes1[..., :2])
  86.     overlap_right_bottom = np.minimum(boxes0[..., 2:], boxes1[..., 2:])
  88.     overlap_area = area_of(overlap_left_top, overlap_right_bottom)
  89.     area0 = area_of(boxes0[..., :2], boxes0[..., 2:])
  90.     area1 = area_of(boxes1[..., :2], boxes1[..., 2:])
  91.     return overlap_area / (area0 + area1 - overlap_area + eps)
  94. def predict(width, height, confidences, boxes, prob_threshold, iou_threshold=0.3, top_k=-1):
  95.     boxes = boxes[0]
  96.     confidences = confidences[0]
  97.     picked_box_probs = []
  98.     picked_labels = []
  99.     for class_index in range(1, confidences.shape[1]):
  100.         probs = confidences[:, class_index]
  101.         mask = probs > prob_threshold
  102.         probs = probs[mask]
  103.         if probs.shape[0] == 0:
  104.             continue
  105.         subset_boxes = boxes[mask, :]
  106.         box_probs = np.concatenate(
  107.             [subset_boxes, probs.reshape(-1, 1)], axis=1)
  108.         box_probs = hard_nms(box_probs,
  109.                              iou_threshold=iou_threshold,
  110.                              top_k=top_k,
  111.                              )
  112.         picked_box_probs.append(box_probs)
  113.         picked_labels.extend([class_index] * box_probs.shape[0])
  114.     if not picked_box_probs:
  115.         return np.array([]), np.array([]), np.array([])
  116.     picked_box_probs = np.concatenate(picked_box_probs)
  117.     picked_box_probs[:, 0] *= width
  118.     picked_box_probs[:, 1] *= height
  119.     picked_box_probs[:, 2] *= width
  120.     picked_box_probs[:, 3] *= height
  121.     return picked_box_probs[:, :4].astype(np.int32), np.array(picked_labels), picked_box_probs[:, 4]
  124. def convert_locations_to_boxes(locations, priors, center_variance,
  125.                                size_variance):
  126.     if len(priors.shape) + 1 == len(locations.shape):
  127.         priors = np.expand_dims(priors, 0)
  128.     return np.concatenate([
  129.         locations[..., :2] * center_variance *
  130.         priors[..., 2:] + priors[..., :2],
  131.         np.exp(locations[..., 2:] * size_variance) * priors[..., 2:]
  132.     ], axis=len(locations.shape) - 1)
  135. def center_form_to_corner_form(locations):
  136.     return np.concatenate([locations[..., :2] - locations[..., 2:] / 2,
  137.                            locations[..., :2] + locations[..., 2:] / 2], len(locations.shape) - 1)
  140. class UltraFaceInference:
  141.     def __init__(self):
  142.         # 模型路径
  143.         model_path = os.path.expanduser(
  144.             '~') + '/Lepi_Data/ros/ultra_face_inference/models/onnx/version-RFB-320_simplified.onnx'
  145.         # opencv dnn 模块加载onnx模型
  146.         self.net = cv2.dnn.readNetFromONNX(model_path)  # onnx version
  147.         # net = cv2.dnn.readNetFromCaffe(args.caffe_prototxt_path, args.caffe_model_path)  # caffe model converted from onnx
  148.         # 缩放尺寸，实际按照该尺寸进行检测，降低精度加快速度
  149.         self.setResize(160, 120)
  150.         # 检测阈值
  151.         self.threshold = 0.7
  152.         # 保存检测结果
  153.         self.faceCount = 0
  154.         self.faceDetections = []
  155.         self.faceData = []
  157.     def detect(self, img_ori):
  158.         time_time = time.time()
  159.         rect = cv2.resize(img_ori, (self.resized_witdh, self.resized_height))
  160.         rect = cv2.cvtColor(rect, cv2.COLOR_BGR2RGB)
  161.         # 准备输入数据
  162.         self.net.setInput(cv2.dnn.blobFromImage(
  163.             rect, 1 / image_std, (self.resized_witdh, self.resized_height), 127))
  164.         # 执行检测（前向传播）
  165.         boxes, scores = self.net.forward(["boxes", "scores"])
  166.         boxes = np.expand_dims(np.reshape(boxes, (-1, 4)), axis=0)
  167.         scores = np.expand_dims(np.reshape(scores, (-1, 2)), axis=0)
  168.         boxes = convert_locations_to_boxes(
  169.             boxes, self.priors, center_variance, size_variance)
  170.         boxes = center_form_to_corner_form(boxes)
  171.         boxes, labels, probs = predict(
  172.             img_ori.shape[1], img_ori.shape[0], scores, boxes, self.threshold)
  173.         self.faceDetections = self.toFaceDetections(boxes, probs)
  174.         self.faceData = self.toFaceData(self.faceDetections)
  175.         for i in range(boxes.shape[0]):
  176.             box = boxes[i, :]
  177.         return boxes, labels, probs
  179.     def drawBoxes(self, img_ori, boxes):
  180.         # 绘制人脸框
  181.         for i in range(boxes.shape[0]):
  182.             box = boxes[i, :]
  183.             cv2.rectangle(img_ori, (box[0], box[1]),
  184.                           (box[2], box[3]), (0, 255, 0), 4)
  185.         # print("inference time: {} s".format(round(time.time() - time_time, 4)))
  186.         return img_ori
  188.     def setThreshold(self, threshold):
  189.         if threshold > 1:
  190.             self.threshold = threshold / 100.0
  191.         else:
  192.             self.threshold = threshold
  194.     def setResize(self, width, height):
  195.         self.resized_witdh = width
  196.         self.resized_height = height
  197.         self.priors = define_img_size(
  198.             [self.resized_witdh, self.resized_height])
  200.     def toFaceDetections(self, face_locations=[], face_probs=[]):
  201.         detections = []
  202.         if len(face_probs) > 0 and len(face_probs) == len(face_locations):
  203.             for (left, top, right, bottom), prob in zip(face_locations, face_probs):
  204.                 face_detection = (prob, [
  205.                     left, top, right, bottom])
  206.                 detections.append(face_detection)
  207.         elif len(face_locations) > 0:
  208.             for (left, top, right, bottom) in face_locations:
  209.                 face_detection = (1, [left, top, right, bottom])
  210.                 detections.append(face_detection)
  211.         return detections
  213.     def toFaceData(self, faceDetections=[]):
  214.         data = []
  215.         for _, face in faceDetections:
  216.             x = (face[0]+face[2])/2
  217.             y = (face[1]+face[3])/2
  218.             w = face[2]-face[0]
  219.             h = face[3]-face[1]
  220.             data.append([x, y, w, h])
  221.         return data
  224. if __name__ == '__main__':
  225.     # inference()
  226.     detector = UltraFaceInference()
  227.     cap = cv2.VideoCapture(0)
  228.     while True:
  229.         # img_path = os.path.join(imgs_path, file_path)
  230.         # img_ori = cv2.imread(img_path)
  231.         ret, img_ori = cap.read()
  232.         time_time = time.time()
  233.         boxes, labels, probs = detector.detect(img_ori)
  234.         img_ori = detector.drawBoxes(img_ori, boxes)
  235.         print("inference time: {} s".format(round(time.time() - time_time, 4)))
  236.         cv2.imshow("ultra_face_inference", np.rot90(
  237.             cv2.resize(img_ori, (320, 240))))
  238.         c = cv2.waitKey(2)
  239.         if c == Key_Esc:
  240.             break
  241.     cv2.destroyAllWindows()
  242.     cap.release()