remove q
Showing
1 changed file
with
0 additions
and
220 deletions
q
deleted
100644 → 0
1 | import os | ||
2 | import numpy as np | ||
3 | import MNN | ||
4 | import cv2 | ||
5 | import logging | ||
6 | from retinaface import PriorBox | ||
7 | |||
8 | |||
9 | def py_cpu_nms(dets, thresh): | ||
10 | """Pure Python NMS baseline.""" | ||
11 | x1 = dets[:, 0] | ||
12 | y1 = dets[:, 1] | ||
13 | x2 = dets[:, 2] | ||
14 | y2 = dets[:, 3] | ||
15 | scores = dets[:, 4] | ||
16 | |||
17 | areas = (x2 - x1 + 1) * (y2 - y1 + 1) | ||
18 | order = scores.argsort()[::-1] | ||
19 | |||
20 | keep = [] | ||
21 | while order.size > 0: | ||
22 | i = order[0] | ||
23 | keep.append(i) | ||
24 | xx1 = np.maximum(x1[i], x1[order[1:]]) | ||
25 | yy1 = np.maximum(y1[i], y1[order[1:]]) | ||
26 | xx2 = np.minimum(x2[i], x2[order[1:]]) | ||
27 | yy2 = np.minimum(y2[i], y2[order[1:]]) | ||
28 | |||
29 | w = np.maximum(0.0, xx2 - xx1 + 1) | ||
30 | h = np.maximum(0.0, yy2 - yy1 + 1) | ||
31 | inter = w * h | ||
32 | ovr = inter / (areas[i] + areas[order[1:]] - inter) | ||
33 | |||
34 | inds = np.where(ovr <= thresh)[0] | ||
35 | order = order[inds + 1] | ||
36 | |||
37 | return keep | ||
38 | |||
39 | |||
40 | def decode_landm(pre, priors, variances): | ||
41 | |||
42 | landms = np.concatenate((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:], | ||
43 | priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:], | ||
44 | priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:], | ||
45 | priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:], | ||
46 | priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:]), 1) | ||
47 | return landms | ||
48 | |||
49 | |||
50 | def decode(loc, priors, variances): | ||
51 | |||
52 | boxes = np.concatenate(( | ||
53 | priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:], | ||
54 | priors[:, 2:] * np.exp(loc[:, 2:] * variances[1])), 1) | ||
55 | boxes[:, :2] -= boxes[:, 2:] / 2 | ||
56 | boxes[:, 2:] += boxes[:, :2] | ||
57 | return boxes | ||
58 | |||
59 | |||
60 | class Face_Detector(object): | ||
61 | def __init__(self, model_path): | ||
62 | logging.info('******** Start Init Face Detector ********') | ||
63 | self.det_interpreter = MNN.Interpreter(model_path) | ||
64 | self.det_session = self.det_interpreter.createSession() | ||
65 | self.det_input_tensor = self.det_interpreter.getSessionInput(self.det_session) | ||
66 | logging.info('******** Success Init Face Detector ********') | ||
67 | |||
68 | def detect(self, frame, thr): | ||
69 | logging.info('******** Start Face Detect ********') | ||
70 | input_size = 320 | ||
71 | img = cv2.resize(frame, (input_size, input_size)) | ||
72 | |||
73 | img = np.float32(img) | ||
74 | im_height, im_width, _ = img.shape | ||
75 | scale = np.array([img.shape[1], img.shape[0], img.shape[1], img.shape[0]]) | ||
76 | scale1 = np.array([img.shape[1], img.shape[0], img.shape[1], img.shape[0], | ||
77 | img.shape[1], img.shape[0], img.shape[1], img.shape[0], | ||
78 | img.shape[1], img.shape[0]]) | ||
79 | |||
80 | w_r = input_size/frame.shape[1] | ||
81 | h_r = input_size/frame.shape[0] | ||
82 | |||
83 | confidence_threshold = 0.02 | ||
84 | vis_threshold = 0.5 | ||
85 | nms_threshold = 0.4 | ||
86 | keep_top_k = 100 | ||
87 | variance = [0.1, 0.2] | ||
88 | img -= (104, 117, 123) | ||
89 | img = img.transpose(2, 0, 1) | ||
90 | img = np.expand_dims(img, axis=0) | ||
91 | |||
92 | input_tensor = MNN.Tensor((1, 3, input_size, input_size), MNN.Halide_Type_Float, img, MNN.Tensor_DimensionType_Caffe) | ||
93 | self.det_input_tensor.copyFrom(input_tensor) | ||
94 | self.det_interpreter.runSession(self.det_session) | ||
95 | |||
96 | bbox_output_tensor = self.det_interpreter.getSessionOutput(self.det_session, 'output0') | ||
97 | conf_output_tensor = self.det_interpreter.getSessionOutput(self.det_session, 'output1') | ||
98 | landmark_output_tensor = self.det_interpreter.getSessionOutput(self.det_session, 'output2') | ||
99 | |||
100 | bbox_output = bbox_output_tensor.getData() | ||
101 | conf_output = conf_output_tensor.getData() | ||
102 | landmark_output = landmark_output_tensor.getData() | ||
103 | |||
104 | norm_confs = list() | ||
105 | for i in range(int(len(conf_output)/2)): | ||
106 | norm_confs.append([conf_output[i * 2 + 0], conf_output[i * 2 + 1]]) | ||
107 | |||
108 | norm_bboxes = list() | ||
109 | for i in range(int(len(conf_output)/2)): | ||
110 | norm_bboxes.append([bbox_output[i * 4 + 0], bbox_output[i * 4 + 1], bbox_output[i * 4 + 2], bbox_output[i * 4 + 3]]) | ||
111 | |||
112 | norm_landmarks = list() | ||
113 | for i in range(int(len(conf_output)/2)): | ||
114 | norm_landmarks.append([landmark_output[i * 10 + 0], landmark_output[i * 10 + 1], | ||
115 | landmark_output[i * 10 + 2], landmark_output[i * 10 + 3], | ||
116 | landmark_output[i * 10 + 4], landmark_output[i * 10 + 5], | ||
117 | landmark_output[i * 10 + 6], landmark_output[i * 10 + 7], | ||
118 | landmark_output[i * 10 + 8], landmark_output[i * 10 + 9]]) | ||
119 | |||
120 | norm_confs = np.array(norm_confs) | ||
121 | norm_bboxes = np.array(norm_bboxes) | ||
122 | norm_landmarks = np.array(norm_landmarks) | ||
123 | |||
124 | priorbox = PriorBox(image_size=(im_height, im_width)) | ||
125 | priors = priorbox.forward() | ||
126 | |||
127 | scores = norm_confs[:, 1] | ||
128 | |||
129 | boxes = decode(norm_bboxes, priors, variance) | ||
130 | boxes = boxes * scale | ||
131 | |||
132 | landms = decode_landm(norm_landmarks, priors, variance) | ||
133 | landms = landms * scale1 | ||
134 | |||
135 | # ignore low scores | ||
136 | inds = np.where(scores > confidence_threshold)[0] | ||
137 | boxes = boxes[inds] | ||
138 | landms = landms[inds] | ||
139 | scores = scores[inds] | ||
140 | |||
141 | # keep top-K before NMS | ||
142 | order = scores.argsort()[::-1][:keep_top_k] | ||
143 | boxes = boxes[order] | ||
144 | landms = landms[order] | ||
145 | scores = scores[order] | ||
146 | |||
147 | # do NMS | ||
148 | dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False) | ||
149 | keep = py_cpu_nms(dets, nms_threshold) | ||
150 | |||
151 | dets = dets[keep, :] | ||
152 | landms = landms[keep] | ||
153 | |||
154 | # keep top-K faster NMS | ||
155 | dets = dets[:keep_top_k, :] | ||
156 | landms = landms[:keep_top_k, :] | ||
157 | |||
158 | dets = np.concatenate((dets, landms), axis=1) | ||
159 | face_bboxes = [] | ||
160 | face_landmarks = [] | ||
161 | max_area = float('-inf') | ||
162 | max_index = 0 | ||
163 | |||
164 | i = 0 | ||
165 | for b in dets: | ||
166 | if b[4] < vis_threshold: | ||
167 | continue | ||
168 | resize_b = [] | ||
169 | x1 = int(b[0] / w_r) | ||
170 | y1 = int(b[1] / h_r) | ||
171 | x2 = int(b[2] / w_r) | ||
172 | y2 = int(b[3] / h_r) | ||
173 | x3 = int(b[5] / w_r) | ||
174 | y3 = int(b[6] / h_r) | ||
175 | x4 = int(b[7] / w_r) | ||
176 | y4 = int(b[8] / h_r) | ||
177 | x5 = int(b[9] / w_r) | ||
178 | y5 = int(b[10] / h_r) | ||
179 | x6 = int(b[11] / w_r) | ||
180 | y6 = int(b[12] / h_r) | ||
181 | x7 = int(b[13] / w_r) | ||
182 | y7 = int(b[14] / h_r) | ||
183 | resize_b = [x1, y1, x2, y2, 0, x3, y3, x4, y4, x5, y5, x6, y6, x7, y7] | ||
184 | |||
185 | # cv2.rectangle(frame, (resize_b[0], resize_b[1]), (resize_b[2], resize_b[3]), (0, 0, 255), 2) | ||
186 | # cv2.circle(frame, (resize_b[5], resize_b[6]), 1, (0, 0, 255), 4) | ||
187 | # cv2.circle(frame, (resize_b[7], resize_b[8]), 1, (0, 255, 255), 4) | ||
188 | # cv2.circle(frame, (resize_b[9], resize_b[10]), 1, (255, 0, 255), 4) | ||
189 | # cv2.circle(frame, (resize_b[11], resize_b[12]), 1, (0, 255, 0), 4) | ||
190 | # cv2.circle(frame, (resize_b[13], resize_b[14]), 1, (255, 0, 0), 4) | ||
191 | |||
192 | area = (resize_b[2] - resize_b[0]) * (resize_b[3] - resize_b[1]) | ||
193 | if area > max_area: | ||
194 | max_area = area | ||
195 | max_index = i | ||
196 | i += 1 | ||
197 | |||
198 | face_bboxes.append([resize_b[0], resize_b[1], resize_b[2], resize_b[3]]) | ||
199 | face_landmarks.append([(resize_b[5], resize_b[6]), | ||
200 | (resize_b[7], resize_b[8]), | ||
201 | (resize_b[9], resize_b[10]), | ||
202 | (resize_b[11], resize_b[12]), | ||
203 | (resize_b[13], resize_b[14])]) | ||
204 | |||
205 | # import time | ||
206 | # cv2.imwrite('results/0.jpg', frame) | ||
207 | return face_bboxes, face_landmarks, max_index | ||
208 | |||
209 | |||
210 | if __name__ == '__main__': | ||
211 | det_face_model_path = r'/home/jwq/PycharmProjects/situ/src/face_det/Pytorch_Retinaface/weights/mobilenet_0.25.mnn' | ||
212 | image_path = r'input/3.jpg' | ||
213 | image_save_path = r'results/3.jpg' | ||
214 | thr = 0.5 | ||
215 | |||
216 | face_detector = Face_Detector(det_face_model_path) | ||
217 | image = cv2.imread(image_path) | ||
218 | face_detector.detect(image, thr) | ||
219 | # image_ploted = face_detector.plot(image, face_bboxes) | ||
220 | # cv2.imwrite(image_save_path, image_ploted) |
-
Please register or sign in to post a comment