Added basic support for toaster level tracking model.

facetracker.py

@@ -18,7 +18,7 @@ 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.65)
+parser.add_argument("-t", "--threshold", type=float, help="Set minimum confidence threshold for face detection", default=None)
 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)
@@ -32,7 +32,7 @@ parser.add_argument("--video-out", help="Set this to the filename of an AVI file
 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("--log-output", help="You can set a filename to console output will be logged here", default="")
-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. Models 1 and 0 tend to be too rigid for expression and blink detection.", default=3, choices=[0, 1, 2, 3])
+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. Models 1 and 0 tend to be too rigid for expression and blink detection.", default=3, choices=[-1, 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("--gaze-tracking", type=int, help="When set to 1, experimental blink detection and gaze tracking are enabled, which makes things slightly slower", default=1)
 parser.add_argument("--face-id-offset", type=int, help="When set, this offset is added to all face ids, which can be useful for mixing tracking data from multiple network sources", default=0)

model.py

@@ -177,6 +177,48 @@ class OpenSeeFaceLandmarks(geffnet.mobilenetv3.MobileNetV3):
     def forward(self, x):
         return self._forward_impl(x)
 
+# lm_modelT for 56x56 30 point inference
+class OpenSeeFaceLandmarks30Pt(geffnet.mobilenetv3.MobileNetV3):
+    def __init__(self, size="large", channel_multiplier=1.0, inference=False):
+        kwargs = geffnet.mobilenetv3._gen_mobilenet_v3([size], channel_multiplier=channel_multiplier)
+        super(OpenSeeFaceLandmarks30Pt, self).__init__(**kwargs)
+        self.up1 = UNetUp(960, 112, 256, (4,4))
+        self.up2 = UNetUp(256, 40, 180, (7,7))
+        self.group = DSConv2d(180, 90, kernels_per_layer=4, groups=3)
+        self.inference = inference
+    def _forward_impl(self, x):
+        x = self.conv_stem(x)
+        x = self.bn1(x)
+        x = self.act1(x)
+        r2 = None
+        r3 = None
+        for i, feature in enumerate(self.blocks):
+            x = feature(x)
+            if i == 4:
+                r3 = x
+            if i == 2:
+                r2 = x
+        x = self.up1(x, r3)
+        x = self.up2(x, r2)
+        x = self.group(x)
+
+        if self.inference:
+            t_main = x[0:30].reshape((30, 7*7))
+            t_m = t_main.argmax(dim=1)
+            indices = t_m.unsqueeze(1)
+            t_conf = t_main.gather(1, indices).squeeze(1)
+            t_off_x = x[30:60].reshape((60, 7*7)).gather(1, indices).squeeze(1)
+            t_off_y = x[60:90].reshape((60, 7*7)).gather(1, indices).squeeze(1)
+            t_off_x = 55. * logit_arr(t_off_x)
+            t_off_y = 55. * logit_arr(t_off_y)
+            t_x = 55. * (t_m / 7.).floor() / 6. + t_off_x
+            t_y = 55. * t_m.remainder(7.).float() / 6. + t_off_y
+            x = (t_conf.mean(), torch.stack([t_x, t_y, t_conf], 1))
+
+        return x
+    def forward(self, x):
+        return self._forward_impl(x)
+
 # Adaptive Wing Loss with offset layers and emphasis for eyes and eyebrows with 66 landmark points
 def AdapWingLoss(pre_hm, gt_hm):
     # pre_hm = pre_hm.to('cpu')
@@ -262,3 +304,7 @@ if __name__== "__main__":
     m=OpenSeeFaceLandmarks("large", 1.0)
     ckpt = torch.load("lm_model3.pth")
     m.load_state_dict(ckpt)
+    print("Checking lm_modelT model")
+    m=OpenSeeFaceLandmarks("large", 1.0)
+    ckpt = torch.load("lm_modelT.pth")
+    m.load_state_dict(ckpt)

tracker.py

@@ -255,7 +255,7 @@ class FeatureExtractor():
             aligned_pts.append(np.array(rotate(a, pt, alpha)))
         return alpha, np.array(aligned_pts)
 
-    def update(self, pts):
+    def update(self, pts, full=True):
         features = {}
         now = time.perf_counter()
 
@@ -270,36 +270,50 @@ class FeatureExtractor():
         f = np.clip((np.mean([f_pts[0,1], f_pts[1,1]]) - np.mean([f_pts[2,1], f_pts[3,1]])) / norm_distance_y, 0, None)
         features["eye_r"] = self.eye_r.update(f, now)
 
-        a3, _ = self.align_points(pts[0], pts[16], [])
+        if full:
-        a4, _ = self.align_points(pts[31], pts[35], [])
+            a3, _ = self.align_points(pts[0], pts[16], [])
-        norm_angle = np.mean(list(map(np.rad2deg, [a1, a2, a3, a4])))
+            a4, _ = self.align_points(pts[31], pts[35], [])
+            norm_angle = np.mean(list(map(np.rad2deg, [a1, a2, a3, a4])))
+
+            a, f_pts = self.align_points(pts[22], pts[26], pts[[22, 23, 24, 25, 26]])
+            features["eyebrow_steepness_l"] = self.eyebrow_steepness_l.update(-np.rad2deg(a) - norm_angle, now)
+            f = np.max(np.abs(np.array(f_pts[1:4]) - f_pts[0, 1])) / norm_distance_y
+            features["eyebrow_quirk_l"] = self.eyebrow_quirk_l.update(f, now)
+
+            a, f_pts = self.align_points(pts[17], pts[21], pts[[17, 18, 19, 20, 21]])
+            features["eyebrow_steepness_r"] = self.eyebrow_steepness_r.update(np.rad2deg(a) - norm_angle, now)
+            f = np.max(np.abs(np.array(f_pts[1:4]) - f_pts[0, 1])) / norm_distance_y
+            features["eyebrow_quirk_r"] = self.eyebrow_quirk_r.update(f, now)
+        else:
+            features["eyebrow_steepness_l"] = 0.
+            features["eyebrow_steepness_r"] = 0.
+            features["eyebrow_quirk_l"] = 0.
+            features["eyebrow_quirk_r"] = 0.
 
-        a, f_pts = self.align_points(pts[22], pts[26], pts[[22, 23, 24, 25, 26]])
-        features["eyebrow_steepness_l"] = self.eyebrow_steepness_l.update(-np.rad2deg(a) - norm_angle, now)
         f = (np.mean([pts[22, 1], pts[26, 1]]) - pts[27, 1]) / norm_distance_y
         features["eyebrow_updown_l"] = self.eyebrow_updown_l.update(f, now)
-        f = np.max(np.abs(np.array(f_pts[1:4]) - f_pts[0, 1])) / norm_distance_y
-        features["eyebrow_quirk_l"] = self.eyebrow_quirk_l.update(f, now)
 
-        a, f_pts = self.align_points(pts[17], pts[21], pts[[17, 18, 19, 20, 21]])
-        features["eyebrow_steepness_r"] = self.eyebrow_steepness_r.update(np.rad2deg(a) - norm_angle, now)
         f = (np.mean([pts[17, 1], pts[21, 1]]) - pts[27, 1]) / norm_distance_y
         features["eyebrow_updown_r"] = self.eyebrow_updown_r.update(f, now)
-        f = np.max(np.abs(np.array(f_pts[1:4]) - f_pts[0, 1])) / norm_distance_y
-        features["eyebrow_quirk_r"] = self.eyebrow_quirk_r.update(f, now)
 
         upper_mouth_line = np.mean([pts[49, 1], pts[50, 1], pts[51, 1]])
         center_line = np.mean([pts[50, 0], pts[60, 0], pts[27, 0], pts[30, 0], pts[64, 0], pts[55, 0]])
 
         f = (upper_mouth_line - pts[62, 1]) / norm_distance_y
         features["mouth_corner_updown_l"] = self.mouth_corner_updown_l.update(f, now)
-        f = abs(center_line - pts[62, 0]) / norm_distance_x
+        if full:
-        features["mouth_corner_inout_l"] = self.mouth_corner_inout_l.update(f, now)
+            f = abs(center_line - pts[62, 0]) / norm_distance_x
+            features["mouth_corner_inout_l"] = self.mouth_corner_inout_l.update(f, now)
+        else:
+            features["mouth_corner_inout_l"] = 0.
 
         f = (upper_mouth_line - pts[58, 1]) / norm_distance_y
         features["mouth_corner_updown_r"] = self.mouth_corner_updown_r.update(f, now)
-        f = abs(center_line - pts[58, 0]) / norm_distance_x
+        if full:
-        features["mouth_corner_inout_r"] = self.mouth_corner_inout_r.update(f, now)
+            f = abs(center_line - pts[58, 0]) / norm_distance_x
+            features["mouth_corner_inout_r"] = self.mouth_corner_inout_r.update(f, now)
+        else:
+            features["mouth_corner_inout_r"] = 0.
 
         f = (np.mean([pts[59, 1], pts[60, 1], pts[61, 1]]) - np.mean([pts[65, 1], pts[64, 1], pts[63, 1]])) / norm_distance_y
         features["mouth_open"] = self.mouth_open.update(f, now)
@@ -316,6 +330,8 @@ class FaceInfo():
         self.tracker = tracker
         self.contour_pts = [0,1,8,15,16,27,28,29,30,31,32,33,34,35,36,39,42,45]
         self.face_3d = copy.copy(self.tracker.face_3d)
+        if self.tracker.model_type < 0:
+            self.contour_pts = [0,2,8,14,16,27,30,33]
         self.reset()
         self.alive = False
         self.coord = None
@@ -385,7 +401,7 @@ class FaceInfo():
         if self.conf < 0.4 or self.pnp_error > 300:
             return
 
-        if not self.tracker.static_model:
+        if self.tracker.model_type > -1 and not self.tracker.static_model:
             max_runs = 1
             eligible = np.delete(np.arange(0, 66), [30])
             changed_any = False
@@ -457,14 +473,19 @@ class FaceInfo():
                 self.update_contour()
 
         self.pts_3d = self.normalize_pts3d(self.pts_3d)
-        self.current_features = self.features.update(self.pts_3d[:, 0:2])
+        if self.tracker.feature_level == 2:
-        self.eye_blink = []
+            self.current_features = self.features.update(self.pts_3d[:, 0:2])
-        self.eye_blink.append(1 - min(max(0, -self.current_features["eye_r"]), 1))
+            self.eye_blink = []
-        self.eye_blink.append(1 - min(max(0, -self.current_features["eye_l"]), 1))
+            self.eye_blink.append(1 - min(max(0, -self.current_features["eye_r"]), 1))
-
+            self.eye_blink.append(1 - min(max(0, -self.current_features["eye_l"]), 1))
+        elif self.tracker.feature_level == 1:
+            self.current_features = self.features.update(self.pts_3d[:, 0:2], False)
+            self.eye_blink = []
+            self.eye_blink.append(1 - min(max(0, -self.current_features["eye_r"]), 1))
+            self.eye_blink.append(1 - min(max(0, -self.current_features["eye_l"]), 1))
 
 class Tracker():
-    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, max_feature_updates=0, static_model=False):
+    def __init__(self, width, height, model_type=3, threshold=None, 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, max_feature_updates=0, static_model=False, feature_level=2):
         options = onnxruntime.SessionOptions()
         options.inter_op_num_threads = 1
         options.intra_op_num_threads = max(max_threads,4)
@@ -478,7 +499,9 @@ class Tracker():
             "lm_model2_opt.onnx",
             "lm_model3_opt.onnx"
         ]
-        model = self.models[self.model_type]
+        model = "lm_modelT_opt.onnx"
+        if model_type >= 0:
+            model = self.models[self.model_type]
         model_base_path = resolve(os.path.join("models"))
         if model_dir is None:
             if not os.path.exists(model_base_path):
@@ -486,6 +509,11 @@ class Tracker():
         else:
             model_base_path = model_dir
 
+        if threshold is None:
+            threshold = 0.6
+            if model_type < 0:
+                threshold = 0.84
+
         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
@@ -529,6 +557,13 @@ class Tracker():
         self.mean = self.mean / self.std
         self.std = self.std * 255.0
 
+        self.mean = - self.mean
+        self.std = 1.0 / self.std
+        self.mean_32 = np.tile(self.mean, [32, 32, 1])
+        self.std_32 = np.tile(self.std, [32, 32, 1])
+        self.mean_224 = np.tile(self.mean, [224, 224, 1])
+        self.std_224 = np.tile(self.std, [224, 224, 1])
+
         # PnP solving
         self.face_3d = np.array([
             [ 0.4551769692672  ,  0.300895790030204, -0.764429433974752],
@@ -623,12 +658,16 @@ class Tracker():
         self.silent = silent
 
         self.res = 224.
+        self.mean_res = self.mean_224
+        self.std_res = self.std_224
         if model_type < 0:
             self.res = 56.
+            self.mean_res = np.tile(self.mean, [56, 56, 1])
+            self.std_res = np.tile(self.std, [56, 56, 1])
         self.res_i = int(self.res)
         self.out_res = 27.
         if model_type < 0:
-            self.out_res = 13.
+            self.out_res = 6.
         self.out_res_i = int(self.out_res) + 1
         self.logit_factor = 16.
         if model_type < 0:
@@ -636,13 +675,16 @@ class Tracker():
 
         self.no_gaze = no_gaze
         self.debug_gaze = False
+        self.feature_level = feature_level
+        if model_type < 0:
+            self.feature_level = min(feature_level, 1)
         self.max_feature_updates = max_feature_updates
         self.static_model = static_model
         self.face_info = [FaceInfo(id, self) for id in range(max_faces)]
         self.fail_count = 0
 
     def detect_faces(self, frame):
-        im = cv2.resize(frame, (224, 224), interpolation=cv2.INTER_LINEAR)[:,:,::-1] / self.std - self.mean
+        im = cv2.resize(frame, (224, 224), interpolation=cv2.INTER_LINEAR)[:,:,::-1] * self.std_224 + self.mean_224
         im = np.expand_dims(im, 0)
         im = np.transpose(im, (0,3,1,2))
         outputs, maxpool = self.detection.run([], {'input': im})
@@ -671,19 +713,27 @@ class Tracker():
         crop_x1, crop_y1, scale_x, scale_y, _ = crop_info
         avg_conf = 0
         res = self.res - 1
-        t_main = tensor[0:66].reshape((66,self.out_res_i * self.out_res_i))
+        c0, c1, c2 = 66, 132, 198
+        if self.model_type < 0:
+            c0, c1, c2 = 30, 60, 90
+        t_main = tensor[0:c0].reshape((c0,self.out_res_i * self.out_res_i))
         t_m = t_main.argmax(1)
         indices = np.expand_dims(t_m, 1)
-        t_conf = np.take_along_axis(t_main, indices, 1).reshape((66,))
+        t_conf = np.take_along_axis(t_main, indices, 1).reshape((c0,))
-        t_off_x = np.take_along_axis(tensor[66:132].reshape((66,self.out_res_i * self.out_res_i)), indices, 1).reshape((66,))
+        t_off_x = np.take_along_axis(tensor[c0:c1].reshape((c0,self.out_res_i * self.out_res_i)), indices, 1).reshape((c0,))
-        t_off_y = np.take_along_axis(tensor[132:198].reshape((66,self.out_res_i * self.out_res_i)), indices, 1).reshape((66,))
+        t_off_y = np.take_along_axis(tensor[c1:c2].reshape((c0,self.out_res_i * self.out_res_i)), indices, 1).reshape((c0,))
-        t_off_x = np.floor(res * logit_arr(t_off_x, self.logit_factor) + 0.5)
+        t_off_x = res * logit_arr(t_off_x, self.logit_factor)
-        t_off_y = np.floor(res * logit_arr(t_off_y, self.logit_factor) + 0.5)
+        t_off_y = res * logit_arr(t_off_y, self.logit_factor)
         t_x = crop_y1 + scale_y * (res * np.floor(t_m / self.out_res_i) / self.out_res + t_off_x)
         t_y = crop_x1 + scale_x * (res * np.floor(np.mod(t_m, self.out_res_i)) / self.out_res + t_off_y)
         avg_conf = np.average(t_conf)
         lms = np.stack([t_x, t_y, t_conf], 1)
         lms[np.isnan(lms).any(axis=1)] = np.array([0.,0.,0.], dtype=np.float32)
+        if self.model_type < 0:
+            lms = lms[[0,0,1,1,1,2,2,2,3,3,3,4,4,4,5,5,6,7,7,8,8,9,10,10,11,11,12,21,21,21,22,23,23,23,23,23,13,14,14,15,16,16,17,18,18,19,20,20,24,25,25,25,26,26,27,27,27,24,24,28,28,28,26,29,29,29]]
+            part_avg = np.mean(np.partition(lms[:,2],3)[0:3])
+            if part_avg < 0.65:
+                avg_conf = part_avg
         return (avg_conf, np.array(lms))
 
     def estimate_depth(self, face_info):
@@ -787,9 +837,7 @@ class Tracker():
     def preprocess(self, im, crop):
         x1, y1, x2, y2 = crop
         im = np.float32(im[y1:y2, x1:x2,::-1]) # Crop and BGR to RGB
-        im = cv2.resize(im, (self.res_i, self.res_i), interpolation=cv2.INTER_LINEAR) / self.std - self.mean
+        im = cv2.resize(im, (self.res_i, self.res_i), interpolation=cv2.INTER_LINEAR) * self.std_res + self.mean_res
-        #im = cv2.resize(im, (224, 224), interpolation=cv2.INTER_LINEAR) / 255.0
-        #im = (im - mean) / std
         im = np.expand_dims(im, 0)
         im = np.transpose(im, (0,3,1,2))
         return im
@@ -823,7 +871,7 @@ class Tracker():
                 full_frame[0:32, 0:32] = im
             else:
                 full_frame[0:32, 32:64] = im
-        im = im.astype(np.float32)[:,:,::-1] / self.std - self.mean
+        im = im.astype(np.float32)[:,:,::-1] * self.mean_32 + self.std_32
         im = np.expand_dims(im, 0)
         im = np.transpose(im, (0,3,2,1))
         return im, np.linalg.inv(transform.params)
@@ -909,6 +957,10 @@ class Tracker():
         return eye_state
 
     def assign_face_info(self, results):
+        if self.max_faces == 1 and len(results) == 1:
+            conf, (lms, eye_state), conf_adjust = results[0]
+            self.face_info[0].update((conf - conf_adjust, (lms, eye_state)), np.array(lms)[:, 0:2].mean(0), self.frame_count)
+            return
         result_coords = []
         adjusted_results = []
         for conf, (lms, eye_state), conf_adjust in results: