diff --git a/model.py b/model.py index d982f03..d73b09d 100644 --- a/model.py +++ b/model.py @@ -158,6 +158,19 @@ class OpenSeeFaceLandmarks(geffnet.mobilenetv3.MobileNetV3): x = self.up1(x, r3) x = self.up2(x, r2) x = self.group(x) + + t_main = x[0:66].reshape((66, 28*28)) + 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[66:132].reshape((66, 28*28)).gather(1, indices).squeeze(1) + t_off_y = x[132:198].reshape((66, 28*28)).gather(1, indices).squeeze(1) + t_off_x = (223. * logit_arr(t_off_x) + 0.5).floor() + t_off_y = (223. * logit_arr(t_off_y) + 0.5).floor() + t_x = 223. * (t_m / 28.).floor() / 27. + t_off_x + t_y = 223. * t_m.remainder(28.).float() / 27. + 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) diff --git a/models/lm_model0_opt.onnx b/models/lm_model0_opt.onnx new file mode 100644 index 0000000..c632f5c Binary files /dev/null and b/models/lm_model0_opt.onnx differ diff --git a/models/mnv3_opt_fast.onnx b/models/lm_model1_opt.onnx similarity index 56% rename from models/mnv3_opt_fast.onnx rename to models/lm_model1_opt.onnx index 07c402f..8d5b440 100644 Binary files a/models/mnv3_opt_fast.onnx and b/models/lm_model1_opt.onnx differ diff --git a/models/mnv3_opt_medium.onnx b/models/lm_model2_opt.onnx similarity index 65% rename from models/mnv3_opt_medium.onnx rename to models/lm_model2_opt.onnx index 915f408..283381c 100644 Binary files a/models/mnv3_opt_medium.onnx and b/models/lm_model2_opt.onnx differ diff --git a/models/mnv3_opt_b.onnx b/models/lm_model3_opt.onnx similarity index 72% rename from models/mnv3_opt_b.onnx rename to models/lm_model3_opt.onnx index 968caf0..67a4857 100644 Binary files a/models/mnv3_opt_b.onnx and b/models/lm_model3_opt.onnx differ diff --git a/models/mnv3_opt_very_fast.onnx b/models/mnv3_opt_very_fast.onnx deleted file mode 100644 index e3547be..0000000 Binary files a/models/mnv3_opt_very_fast.onnx and /dev/null differ diff --git a/tracker.py b/tracker.py index 91f0967..8add85f 100644 --- a/tracker.py +++ b/tracker.py @@ -131,8 +131,9 @@ def matrix_to_quaternion(m): return q def worker_thread(session, frame, input, crop_info, queue, input_name, idx, tracker): - output = session.run([], {input_name: input})[0] - conf, lms = tracker.landmarks(output[0], crop_info) + output = session.run([], {input_name: input}) + lms = np.array(output[1]) * np.array([crop_info[3], crop_info[2], 1.]) + np.array([crop_info[1], crop_info[0], 0.]) + conf = output[0] if conf > tracker.threshold: eye_state = tracker.get_eye_state(frame, lms, single=True) queue.put((session, conf, (lms, eye_state), crop_info, idx)) @@ -465,10 +466,10 @@ class Tracker(): options.log_severity_level = 3 self.model_type = model_type self.models = [ - "mnv3_opt_very_fast.onnx", - "mnv3_opt_fast.onnx", - "mnv3_opt_medium.onnx", - "mnv3_opt_b.onnx" + "lm_model0_opt.onnx", + "lm_model1_opt.onnx", + "lm_model2_opt.onnx", + "lm_model3_opt.onnx" ] model = self.models[self.model_type] model_base_path = resolve(os.path.join("models")) @@ -648,27 +649,6 @@ class Tracker(): results[:, [1,3]] *= frame.shape[0] / 224. return results - def landmarks(self, tensor, crop_info): - crop_x1, crop_y1, scale_x, scale_y, _ = crop_info - avg_conf = 0 - lms = [] - res = self.res - 1 - 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 = 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 * (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, np.array(lms)) - def estimate_depth(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) @@ -1000,8 +980,9 @@ class Tracker(): start_model = time.perf_counter() outputs = {} if num_crops == 1: - output = self.session.run([], {self.input_name: crops[0]})[0] - conf, lms = self.landmarks(output[0], crop_info[0]) + output = self.session.run([], {self.input_name: crops[0]}) + lms = np.array(output[1]) * np.array([crop_info[0][3], crop_info[0][2], 1.]) + np.array([crop_info[0][1], crop_info[0][0], 0.]) + conf = output[0] if conf > self.threshold: eye_state = self.get_eye_state(frame, lms) outputs[crop_info[0]] = (conf, (lms, eye_state), 0)