The model selection parameter should now be easier to understand and two less useful models have been removed to save space. Width, height and FPS are now written to log file and both width and height are also transmitted over the network. The video output function has been changed to use a lossless codec, but videos need to be converted to a more common format before importing them to Vegas now.

facetracker.py

@@ -25,7 +25,7 @@ parser.add_argument("--discard-after", type=int, help="Set the how long the trac
 parser.add_argument("--video-out", help="Set this to the filename of an AVI file to save the tracking visualization as a video", default=None)
 parser.add_argument("--raw-rgb", type=int, help="When this is set, raw RGB frames of the size given with \"-W\" and \"-H\" are read from standard input instead of reading a video", default=0)
 parser.add_argument("--log-data", help="You can set a filename to which tracking data will be logged here", default="")
-parser.add_argument("--model", type=int, help="This can be used to select between the two models with 0 being the original model that might be overall more accurate, 1 being a model that might be able to fit eye and eyebrow shapes more accurately and 2 being a much faster, but lower accuracy model", default=1, choices=[0, 1, 2])
+parser.add_argument("--model", type=int, help="This can be used to select the tracking model. Higher numbers are models with better tracking quality, but slower speed. The recommended models are 3 and 1.", default=3, choices=[0, 1, 2, 3])
 parser.add_argument("--model-dir", help="This can be used to specify the path to the directory containing the .onnx model files", default=None)
 parser.add_argument("--high-quality-3d", type=int, help="When set to 1, more nose points are used when estimating the face pose", default=1)
 parser.add_argument("--gaze-tracking", type=int, help="When set to 1, experimental blink detection and gaze tracking are enabled, which makes things slightly slower", default=1)
@@ -79,7 +79,7 @@ frame_count = 0
 
 if args.log_data != "":
     log = open(args.log_data, "w")
-    log.write("Frame,Time,Face,RightOpen,LeftOpen,AverageConfidence,Success3D,PnPError,RotationQuat.X,RotationQuat.Y,RotationQuat.Z,RotationQuat.W,RVec.X,RVec.Y,RVec.Z,TVec.X,TVec.Y,TVec.Z")
+    log.write("Frame,Time,Width,Height,FPS,Face,RightOpen,LeftOpen,AverageConfidence,Success3D,PnPError,RotationQuat.X,RotationQuat.Y,RotationQuat.Z,RotationQuat.W,Euler.X,Euler.Y,Euler.Z,RVec.X,RVec.Y,RVec.Z,TVec.X,TVec.Y,TVec.Z")
     for i in range(66):
         log.write(f",Landmark[{i}].X,Landmark[{i}].Y,Landmark[{i}].Confidence")
     for i in range(66):
@@ -106,7 +106,7 @@ try:
             sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
             tracker = Tracker(width, height, threshold=args.threshold, max_threads=args.max_threads, max_faces=args.faces, discard_after=args.discard_after, scan_every=args.scan_every, silent=False if args.silent == 0 else True, model_type=args.model, pnp_quality=args.high_quality_3d, model_dir=args.model_dir, no_gaze=False if args.gaze_tracking != 0 else True)
             if not args.video_out is None:
-                out = cv2.VideoWriter(args.video_out, cv2.VideoWriter_fourcc('M','J','P','G'), 24, (width,height))
+                out = cv2.VideoWriter(args.video_out, cv2.VideoWriter_fourcc('F','F','V','1'), 24, (width,height))
 
         inference_start = time.perf_counter()
         faces = tracker.predict(frame)
@@ -124,6 +124,8 @@ try:
             if not success_3d:
                 pts_3d = np.zeros((70, 3), np.float32)
             packet.extend(bytearray(struct.pack("d", now)))
+            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)))
@@ -139,7 +141,7 @@ try:
             packet.extend(bytearray(struct.pack("f", translation[1])))
             packet.extend(bytearray(struct.pack("f", translation[2])))
             if not log is None:
-                log.write(f"{frame_count},{now},{face_num},{right_open},{left_open},{conf},{success_3d},{pnp_error},{quaternion[0]},{quaternion[1]},{quaternion[2]},{quaternion[3]},{rotation[0][0]},{rotation[0][1]},{rotation[0][2]},{translation[0][0]},{translation[0][1]},{translation[0][2]}")
+                log.write(f"{frame_count},{now},{width},{height},{args.fps},{face_num},{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:
                 packet.extend(bytearray(struct.pack("f", c)))
             for pt_num, (x,y,c) in enumerate(lms):

tracker.py

@@ -88,7 +88,7 @@ def worker_thread(session, input, crop_info, queue, input_name, idx, tracker):
     queue.put((session, conf, lms, crop_info, idx))
 
 class Tracker():
-    def __init__(self, width, height, model_type=1, 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):
+    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):
         options = onnxruntime.SessionOptions()
         options.inter_op_num_threads = 1
         options.intra_op_num_threads = max_threads
@@ -96,7 +96,12 @@ class Tracker():
         options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
         options.log_severity_level = 3
         self.model_type = model_type
-        self.models = ["snv2_opt.onnx", "snv2_opt_b.onnx", "snv2_opt_fast.onnx"]
+        self.models = [
+            "snv2_opt_fast.onnx",
+            "mnv3_opt_fast.onnx",
+            "snv2_opt_b.onnx",
+            "mnv3_opt_b.onnx"
+        ]
         model = self.models[self.model_type]
         model_base_path = resolve(os.path.join("models"))
         if model_dir is None:
@@ -243,7 +248,7 @@ class Tracker():
         self.scan_every = scan_every
         self.bbox_growth = bbox_growth
         self.silent = silent
-        self.res = 224. if self.model_type < 2 else 112.
+        self.res = 224. if self.model_type != 0 else 112.
         self.res_i = int(self.res)
         self.no_gaze = no_gaze
 
@@ -337,7 +342,7 @@ class Tracker():
         #im = (im - mean) / std
         im = np.expand_dims(im, 0)
         im = np.transpose(im, (0,3,1,2))
-        if self.model_type > 1:
+        if self.model_type == 0:
             im = im.mean(1)
             im = np.expand_dims(im, 1)
         return im