diff --git a/facetracker.py b/facetracker.py index a11a739..6959815 100644 --- a/facetracker.py +++ b/facetracker.py @@ -16,8 +16,8 @@ else: parser.add_argument("-H", "--height", type=int, help="Set raw RGB height", default=360) parser.add_argument("-c", "--capture", help="Set camera ID (0, 1...) or video file", default="0") parser.add_argument("-m", "--max-threads", type=int, help="Set the maximum number of threads", default=1) -parser.add_argument("-t", "--threshold", type=float, help="Set minimum confidence threshold for face detection", default=0.4) -parser.add_argument("-v", "--visualize", type=int, help="Set this to 1 to visualize the tracking, to 2 to also show face ids or to 3 to add numbers to the point display", default=0) +parser.add_argument("-t", "--threshold", type=float, help="Set minimum confidence threshold for face detection", default=0.6) +parser.add_argument("-v", "--visualize", type=int, help="Set this to 1 to visualize the tracking, to 2 to also show face ids, to 3 to add confidence values or to 4 to add numbers to the point display", default=0) parser.add_argument("-P", "--pnp-points", type=int, help="Set this to 1 to add the 3D fitting points to the visualization", default=0) parser.add_argument("-s", "--silent", type=int, help="Set this to 1 to prevent text output on the console", default=0) parser.add_argument("--faces", type=int, help="Set the maximum number of faces (slow)", default=1) @@ -159,6 +159,10 @@ try: 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))) + if args.visualize > 1: + frame = cv2.putText(frame, str(f.id), (int(f.bbox[0]), int(f.bbox[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255,0,255)) + if args.visualize > 2: + frame = cv2.putText(frame, f"{f.conf:.4f}", (int(f.bbox[0] + 18), int(f.bbox[1] - 6)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0)) for pt_num, (x,y,c) in enumerate(f.lms): packet.extend(bytearray(struct.pack("f", y))) packet.extend(bytearray(struct.pack("f", x))) @@ -171,9 +175,7 @@ try: x = int(x + 0.5) y = int(y + 0.5) if args.visualize != 0 or not out is None: - if args.visualize > 1: - frame = cv2.putText(frame, str(f.id), (int(f.bbox[0]), int(f.bbox[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255,0,255)) - if args.visualize > 2: + if args.visualize > 3: frame = cv2.putText(frame, str(pt_num), (int(y), int(x)), cv2.FONT_HERSHEY_SIMPLEX, 0.25, (255,255,0)) color = (0, 0, 255) if pt_num >= 66: diff --git a/models/mnv3_opt_b.onnx b/models/mnv3_opt_b.onnx index 7bfc796..968caf0 100644 Binary files a/models/mnv3_opt_b.onnx and b/models/mnv3_opt_b.onnx differ diff --git a/run.bat b/run.bat index 2d3aeab..50a004f 100644 --- a/run.bat +++ b/run.bat @@ -8,6 +8,6 @@ set /p width=Select the width: set /p height=Select the height: set /p fps=Select the FPS: -facetracker -c %cameraNum% -W %width% -H %height% -F %fps% -v 1 -P 1 +facetracker -c %cameraNum% -W %width% -H %height% -F %fps% -v 3 -P 1 pause \ No newline at end of file diff --git a/tracker.py b/tracker.py index 810fd89..c89adcc 100644 --- a/tracker.py +++ b/tracker.py @@ -405,10 +405,10 @@ class FaceInfo(): 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, model_dir=None, no_gaze=False, use_retinaface=False): + def __init__(self, width, height, model_type=3, threshold=0.6, max_faces=1, discard_after=5, scan_every=3, bbox_growth=0.0, max_threads=4, silent=False, model_dir=None, no_gaze=False, use_retinaface=False): options = onnxruntime.SessionOptions() options.inter_op_num_threads = 1 - options.intra_op_num_threads = max_threads + options.intra_op_num_threads = max(max_threads,4) options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL options.log_severity_level = 3 @@ -427,7 +427,7 @@ class Tracker(): else: model_base_path = model_dir - self.retinaface = RetinaFaceDetector(model_path=os.path.join(model_base_path, "retinaface_640x640_opt.onnx"), json_path=os.path.join(model_base_path, "priorbox_640x640.json"), threads=max_threads, top_k=max_faces, res=(640, 640)) + self.retinaface = RetinaFaceDetector(model_path=os.path.join(model_base_path, "retinaface_640x640_opt.onnx"), json_path=os.path.join(model_base_path, "priorbox_640x640.json"), threads=max(max_threads,4), top_k=max_faces, res=(640, 640)) self.retinaface_scan = RetinaFaceDetector(model_path=os.path.join(model_base_path, "retinaface_640x640_opt.onnx"), json_path=os.path.join(model_base_path, "priorbox_640x640.json"), threads=2, top_k=max_faces, res=(640, 640)) self.use_retinaface = use_retinaface @@ -453,7 +453,7 @@ class Tracker(): options = onnxruntime.SessionOptions() options.inter_op_num_threads = 1 - options.intra_op_num_threads = max_threads + options.intra_op_num_threads = max(max_threads,4) options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL options.log_severity_level = 3 @@ -597,18 +597,19 @@ class Tracker(): crop_x1, crop_y1, scale_x, scale_y, _ = crop_info avg_conf = 0 lms = [] + res = self.res - 1 if self.model_type == 3 else self.res for i in range(0, 66): m = int(tensor[i].argmax()) x = m // 28 y = m % 28 conf = float(tensor[i][x,y]) avg_conf = avg_conf + conf - off_x = self.res * ((1. * logit(tensor[66 + i][x, y])) - 0.0) - off_y = self.res * ((1. * logit(tensor[66 * 2 + i][x, y])) - 0.0) + off_x = res * ((1. * logit(tensor[66 + i][x, y])) - 0.0) + off_y = res * ((1. * logit(tensor[66 * 2 + i][x, y])) - 0.0) off_x = math.floor(off_x + 0.5) off_y = math.floor(off_y + 0.5) - lm_x = crop_y1 + scale_y * (self.res * (float(x) / 28.) + off_x) - lm_y = crop_x1 + scale_x * (self.res * (float(y) / 28.) + off_y) + lm_x = crop_y1 + scale_y * (res * (float(x) / 27.) + off_x) + lm_y = crop_x1 + scale_x * (res * (float(y) / 27.) + off_y) lms.append((lm_x,lm_y,conf)) avg_conf = avg_conf / 66. return (avg_conf, lms)