diff --git a/facetracker.py b/facetracker.py index 62d19d5..e221229 100644 --- a/facetracker.py +++ b/facetracker.py @@ -115,37 +115,37 @@ try: tracking_frames += 1 packet = bytearray() detected = False - for face_num, (conf, lms, success_3d, pnp_error, quaternion, euler, rotation, translation, pts_3d, (right_open, left_open), (face_y, face_x, face_height, face_width), face_id) in enumerate(faces): - left_state = "O" if left_open > 0.5 else "-" - right_state = "O" if right_open > 0.5 else "-" + for face_num, f in enumerate(faces): + right_state = "O" if f.eye_blink[0] > 0.5 else "-" + left_state = "O" if f.eye_blink[1] > 0.5 else "-" if args.silent == 0: - print(f"Confidence[{face_id}]: {conf:.4f} / 3D fitting error: {pnp_error:.4f} / Eyes: {left_state}, {right_state}") + print(f"Confidence[{f.id}]: {f.conf:.4f} / 3D fitting error: {f.pnp_error:.4f} / Eyes: {left_state}, {right_state}") detected = True - if not success_3d: + if not f.success: pts_3d = np.zeros((70, 3), np.float32) packet.extend(bytearray(struct.pack("d", now))) - packet.extend(bytearray(struct.pack("i", face_id))) + packet.extend(bytearray(struct.pack("i", f.id))) packet.extend(bytearray(struct.pack("f", width))) packet.extend(bytearray(struct.pack("f", height))) - packet.extend(bytearray(struct.pack("f", right_open))) - packet.extend(bytearray(struct.pack("f", left_open))) - packet.extend(bytearray(struct.pack("B", 1 if success_3d else 0))) - packet.extend(bytearray(struct.pack("f", pnp_error))) - packet.extend(bytearray(struct.pack("f", quaternion[0]))) - packet.extend(bytearray(struct.pack("f", quaternion[1]))) - packet.extend(bytearray(struct.pack("f", quaternion[2]))) - packet.extend(bytearray(struct.pack("f", quaternion[3]))) - packet.extend(bytearray(struct.pack("f", euler[0]))) - packet.extend(bytearray(struct.pack("f", euler[1]))) - packet.extend(bytearray(struct.pack("f", euler[2]))) - packet.extend(bytearray(struct.pack("f", translation[0]))) - packet.extend(bytearray(struct.pack("f", translation[1]))) - packet.extend(bytearray(struct.pack("f", translation[2]))) + packet.extend(bytearray(struct.pack("f", f.eye_blink[0]))) + packet.extend(bytearray(struct.pack("f", f.eye_blink[1]))) + packet.extend(bytearray(struct.pack("B", 1 if f.success else 0))) + packet.extend(bytearray(struct.pack("f", f.pnp_error))) + packet.extend(bytearray(struct.pack("f", f.quaternion[0]))) + packet.extend(bytearray(struct.pack("f", f.quaternion[1]))) + packet.extend(bytearray(struct.pack("f", f.quaternion[2]))) + packet.extend(bytearray(struct.pack("f", f.quaternion[3]))) + packet.extend(bytearray(struct.pack("f", f.euler[0]))) + packet.extend(bytearray(struct.pack("f", f.euler[1]))) + packet.extend(bytearray(struct.pack("f", f.euler[2]))) + packet.extend(bytearray(struct.pack("f", f.translation[0]))) + packet.extend(bytearray(struct.pack("f", f.translation[1]))) + packet.extend(bytearray(struct.pack("f", f.translation[2]))) if not log is None: - log.write(f"{frame_count},{now},{width},{height},{args.fps},{face_num},{face_id},{right_open},{left_open},{conf},{success_3d},{pnp_error},{quaternion[0]},{quaternion[1]},{quaternion[2]},{quaternion[3]},{euler[0]},{euler[1]},{euler[2]},{rotation[0]},{rotation[1]},{rotation[2]},{translation[0]},{translation[1]},{translation[2]}") - for (x,y,c) in lms: + log.write(f"{frame_count},{now},{width},{height},{args.fps},{face_num},{f.id},{f.eye_blink[0]},{f.eye_blink[1]},{f.conf},{f.success},{f.pnp_error},{f.quaternion[0]},{f.quaternion[1]},{f.quaternion[2]},{f.quaternion[3]},{f.euler[0]},{f.euler[1]},{f.euler[2]},{f.rotation[0]},{f.rotation[1]},{f.rotation[2]},{f.translation[0]},{f.translation[1]},{f.translation[2]}") + for (x,y,c) in f.lms: packet.extend(bytearray(struct.pack("f", c))) - for pt_num, (x,y,c) in enumerate(lms): + for pt_num, (x,y,c) in enumerate(f.lms): packet.extend(bytearray(struct.pack("f", y))) packet.extend(bytearray(struct.pack("f", x))) if not log is None: @@ -169,7 +169,7 @@ try: if not (x < 0 or y < 0 or x >= height or y >= width): frame[int(x), int(y)] = (0, 0, 255) if args.pnp_points != 0 and (args.visualize != 0 or not out is None): - projected = cv2.projectPoints(tracker.contour, rotation, translation, tracker.camera, tracker.dist_coeffs) + projected = cv2.projectPoints(tracker.contour, f.rotation, f.translation, tracker.camera, tracker.dist_coeffs) for [(x,y)] in projected[0]: if not (x < 0 or y < 0 or x >= height or y >= width): frame[int(x), int(y)] = (0, 255, 255) @@ -182,7 +182,7 @@ try: x -= 1 if not (x < 0 or y < 0 or x >= height or y >= width): frame[int(x), int(y)] = (0, 255, 255) - for (x,y,z) in pts_3d: + for (x,y,z) in f.pts_3d: packet.extend(bytearray(struct.pack("f", x))) packet.extend(bytearray(struct.pack("f", -y))) packet.extend(bytearray(struct.pack("f", -z))) diff --git a/tracker.py b/tracker.py index 61a223c..0923fb8 100644 --- a/tracker.py +++ b/tracker.py @@ -46,7 +46,7 @@ def group_rects(rects): rect_groups[name] = [group, -1, []] else: group = rect_groups[name][0] - for j, other_rect in enumerate(rects): + for j, other_rect in enumerate(rects): if i == j: continue; inter = intersects(rect, other_rect) @@ -90,17 +90,34 @@ def worker_thread(session, input, crop_info, queue, input_name, idx, tracker): class FaceInfo(): def __init__(self, id): self.id = id + self.frame_count = -1 + self.reset() + self.alive = False + self.coord = None + + def reset(self): + self.alive = False + self.conf = None + self.lms = None + self.eye_state = None self.rotation = None self.translation = None - self.coord = None - self.frame_count = -1 - def update(self, coord, frame_count): + self.success = None + self.quaternion = None + self.euler = None + self.pnp_error = None + self.pts_3d = None + self.eye_blink = None + self.bbox = None + + def update(self, result, coord, frame_count): self.frame_count = frame_count - if coord is None: - self.rotation = None - self.translation = None + if result is None: + self.reset() else: + self.conf, (self.lms, self.eye_state) = result self.coord = coord + self.alive = True class Tracker(): def __init__(self, width, height, model_type=3, threshold=0.4, max_faces=1, discard_after=5, scan_every=3, bbox_growth=0.0, max_threads=4, silent=False, pnp_quality=1, model_dir=None, no_gaze=False): @@ -124,10 +141,10 @@ class Tracker(): model_base_path = resolve(os.path.join("..", "models")) else: model_base_path = model_dir - + # Single face instance with multiple threads self.session = onnxruntime.InferenceSession(os.path.join(model_base_path, model), sess_options=options) - + # Multiple faces with single threads self.sessions = [] self.max_workers = min(max_threads, max_faces) @@ -144,7 +161,7 @@ class Tracker(): options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL self.sessions.append(onnxruntime.InferenceSession(os.path.join(model_base_path, model), sess_options=options)) self.input_name = self.session.get_inputs()[0].name - + options = onnxruntime.SessionOptions() options.inter_op_num_threads = 1 options.intra_op_num_threads = max_threads @@ -165,8 +182,6 @@ class Tracker(): # PnP solving self.high_quality_3d = False if pnp_quality == 0 else True - self.rotation = None - self.translation = None self.face_3d = np.array([ [0.50194, 0.28030, -0.49561], [0.48377, 0.15031, -0.50000], @@ -240,7 +255,7 @@ class Tracker(): [0.25799, 0.27608, -0.24967], [-0.25799, 0.27608, -0.24967], ], np.float32) * np.array([1.0, 1.0, 1.3]) - + self.contour = np.empty((21 if self.high_quality_3d else 18, 3)) self.contour[0:17] = self.face_3d[0:17] if self.high_quality_3d: @@ -289,8 +304,8 @@ class Tracker(): avg_conf = avg_conf / 66. return (avg_conf, lms) - def estimateDepth(self, lms, eye_state, face_info): - lms = np.concatenate((lms, np.array([[eye_state[0][1], eye_state[0][2], eye_state[0][3]], [eye_state[1][1], eye_state[1][2], eye_state[1][3]]], np.float)), 0) + def estimateDepth(self, face_info): + lms = np.concatenate((face_info.lms, np.array([[face_info.eye_state[0][1], face_info.eye_state[0][2], face_info.eye_state[0][3]], [face_info.eye_state[1][1], face_info.eye_state[1][2], face_info.eye_state[1][3]]], np.float)), 0) image_pts = np.empty((21 if self.high_quality_3d else 18, 2), np.float32) image_pts[0:17] = np.array(lms)[0:17,0:2] @@ -343,7 +358,7 @@ class Tracker(): e2 = lms[i][1] - reference[1] pnp_error += e1*e1 + e2*e2 pt_3d = np.array([lms[i][0] * depth, lms[i][1] * depth, depth], np.float32) - pt_3d = self.inverse_camera.dot(pt_3d) + pt_3d = self.inverse_camera.dot(pt_3d) pt_3d = pt_3d - face_info.translation pt_3d = inverse_rotation.dot(pt_3d) pts_3d[i,:] = pt_3d[:] @@ -402,7 +417,7 @@ class Tracker(): x = m // 8 y = m % 8 conf = float(results[i][0][x,y]) - + off_y = 32.0 * logit(results[i][1][y, x], 8.0) off_x = 32.0 * logit(results[i][2][y, x], 8.0) eye_y = e_y[i] + scale[i] * (32.0 * float(x) / 8.0 + off_x) @@ -411,28 +426,40 @@ class Tracker(): return eye_state - def get_face_info(self, lms): - coord = np.array(lms)[:, 0:2].mean(0) - empty = None - closest = None - best_distance = None + def assign_face_info(self, results): + result_coords = [] + for conf, (lms, eye_state) in results: + result_coords.append(np.array(lms)[:, 0:2].mean(0)) + candidates = [[]] * 4 + max_dist = 2 * np.linalg.norm(np.array([self.width, self.height])) + for i, face_info in enumerate(self.face_info): + for j, coord in enumerate(result_coords): + if face_info.coord is None: + candidates[i].append((max_dist, i, j)) + else: + candidates[i].append((np.linalg.norm(face_info.coord - coord), i, j)) + for i, candidate in enumerate(candidates): + candidates[i] = sorted(candidate) + found = 0 + target = len(results) + used_results = {} + used_faces = {} + while found < target: + min_list = min(candidates) + candidate = min_list.pop(0) + face_idx = candidate[1] + result_idx = candidate[2] + if not result_idx in used_results and not face_idx in used_faces: + self.face_info[face_idx].update(results[result_idx], result_coords[result_idx], self.frame_count) + min_list.clear() + used_results[result_idx] = True + used_faces[face_idx] = True + found += 1 + if len(min_list) == 0: + min_list.append((2 * max_dist, face_idx, result_idx)) for face_info in self.face_info: if face_info.frame_count != self.frame_count: - if face_info.coord is None: - if empty is None: - empty = face_info - else: - d = np.linalg.norm(face_info.coord - coord) - if best_distance is None or d < best_distance: - best_distance = d - closest = face_info - if closest is not None: - closest.update(coord, self.frame_count) - return closest - if empty is not None: - empty.update(coord, self.frame_count) - return empty - return FaceInfo(-1) + face_info.update(None, None, self.frame_count) def predict(self, frame, additional_faces=[]): self.frame_count += 1 @@ -478,7 +505,7 @@ class Tracker(): scale_x = float(crop_x2 - crop_x1) / self.res scale_y = float(crop_y2 - crop_y1) / self.res - + if crop_x2 - crop_x1 < 4 or crop_y2 - crop_y1 < 4: continue @@ -526,27 +553,23 @@ class Tracker(): if conf > self.threshold and best_results[group_id][0] < conf: best_results[group_id][0] = conf best_results[group_id][1] = (lms, eye_state) - - sorted_results = sorted(best_results.values(), key=lambda x: x[0], reverse=True) + + sorted_results = sorted(best_results.values(), key=lambda x: x[0], reverse=True)[:self.max_faces] + self.assign_face_info(sorted_results) duration_model = 1000 * (time.perf_counter() - start_model) results = [] detected = [] - detections = 0 - for conf, (lms, eye_state) in sorted_results: - if detections >= self.max_faces: - break; - if conf > self.threshold: - detections += 1 + for face_info in self.face_info: + if face_info.alive and face_info.conf > self.threshold: start_pnp = time.perf_counter() - face_info = self.get_face_info(lms) - success, quaternion, euler, pnp_error, pts_3d, lms = self.estimateDepth(lms, eye_state, face_info) + face_info.success, face_info.quaternion, face_info.euler, face_info.pnp_error, face_info.pts_3d, face_info.lms = self.estimateDepth(face_info) duration_pnp += 1000 * (time.perf_counter() - start_pnp) min_x = 10000 max_x = -1 min_y = 10000 max_y = -1 - for (x,y,c) in lms[0:66]: + for (x,y,c) in face_info.lms[0:66]: if x > max_x: max_x = x if y > max_y: @@ -557,20 +580,17 @@ class Tracker(): min_y = y w = max_x - min_x y = max_y - min_y - detected.append((min_y, min_x, max_y - min_y, max_x - min_x)) - results.append((conf, lms, success, pnp_error, quaternion, euler, face_info.rotation, face_info.translation, pts_3d, (eye_state[0][0], eye_state[1][0]), (min_y, min_x, max_y - min_y, max_x - min_x), face_info.id)) - - for face_info in self.face_info: - if face_info.frame_count != self.frame_count: - face_info.update(None, self.frame_count) + bbox = (min_y, min_x, max_y - min_y, max_x - min_x) + detected.append(bbox) + face_info.bbox = bbox + face_info.eye_blink = (eye_state[0][0], eye_state[1][0]) + results.append(face_info) if len(detected) > 0: self.detected = len(detected) self.faces = detected self.discard = 0 else: - self.rotation = None - self.translation = None self.detected = 0 self.discard += 1 if self.discard > self.discard_after: @@ -578,7 +598,7 @@ class Tracker(): else: if self.bbox_growth > 0: faces = [] - for (x,y,w,h) in self.faces: + for (x,y,w,h) in self.faces: x -= w * self.bbox_growth y -= h * self.bbox_growth w += 2 * w * self.bbox_growth @@ -589,7 +609,7 @@ class Tracker(): duration = (time.perf_counter() - start) * 1000 if not self.silent: print(f"Took {duration:.2f}ms (detect: {duration_fd:.2f}ms, crop: {duration_pp:.2f}, track: {duration_model:.2f}ms, 3D points: {duration_pnp:.2f}ms)") - - results = sorted(results, key=lambda x: x[11]) + + results = sorted(results, key=lambda x: x.id) return results