diff --git a/eyes.py b/eyes.py index 5e9ee7b..cd723f5 100644 --- a/eyes.py +++ b/eyes.py @@ -149,17 +149,7 @@ class EyeTracker(): self.faceCenter = None self.faceRadius = None - model_base_path = os.path.join(os.path.dirname(__file__), os.path.join("models")) - providersList = onnxruntime.capi._pybind_state.get_available_providers() - options = onnxruntime.SessionOptions() - options.inter_op_num_threads = 1 - options.intra_op_num_threads = 1 - options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL - options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL - options.log_severity_level = 3 - self.gaze_model = onnxruntime.InferenceSession(os.path.join(model_base_path, "mnv3_gaze32_split_opt.onnx"), sess_options=options, providers=providersList) - - def get_eye_state(self,frame, lms): + def get_eye_state(self, models, frame, lms): lms = self.extract_face(frame, np.array(lms)[:,0:2][:,::-1]) @@ -177,7 +167,7 @@ class EyeTracker(): return [[1.,0.,0.,0.],[1.,0.,0.,0.]] #Early exit if one of the eyes doesn't have data both_eyes = np.concatenate((self.rightEye.image, self.leftEye.image)) - self.rightEye.results, self.leftEye.results = self.gaze_model.run([], {"input": both_eyes})[0] + self.rightEye.results, self.leftEye.results = models.gazeTracker.run([], {"input": both_eyes})[0] self.rightEye.calculateEye() self.leftEye.calculateEye() diff --git a/face.py b/face.py index 6549c9a..52152b4 100644 --- a/face.py +++ b/face.py @@ -4,10 +4,8 @@ import featureExtractor import emilianaFeatureExtractor class FaceInfo(): - def __init__(self, id, tracker, featureType): + def __init__(self, featureType): self.featureType = featureType - self.id = id - self.tracker = tracker self.face_3d = np.array([ [ 0.4551769692672 , 0.300895790030204, -0.764429433974752], [ 0.448998827123556, 0.166995837790733, -0.765143004071253], diff --git a/facedetection.py b/facedetection.py index 53b5fc2..909a859 100644 --- a/facedetection.py +++ b/facedetection.py @@ -5,34 +5,9 @@ import onnxruntime import time -#kinda split on how to handle this one -#It's the cause of all of my slowest frames -#but it's not run often -#and I don't know how to make it faster - -class FaceDetectionModel(): - def __init__(self): - model_base_path = os.path.join(os.path.dirname(__file__), os.path.join("models")) - providersList = onnxruntime.capi._pybind_state.get_available_providers() - - options = onnxruntime.SessionOptions() - options.inter_op_num_threads = 1 - options.intra_op_num_threads = 1 - options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL - options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL - options.log_severity_level = 3 - self.detection = onnxruntime.InferenceSession(os.path.join(model_base_path, "mnv3_detection_opt.onnx"), sess_options=options, providers=providersList) - - def run(self, image): - outputs, maxpool = self.detection.run([], {'input': image}) - return outputs[0] - - - class FaceDetector(): def __init__(self, detection_threshold=0.6): - self.model = FaceDetectionModel() self.mean = np.float32(np.array([-2.1179, -2.0357, -1.8044])) self.std = np.float32(np.array([0.0171, 0.0175, 0.0174])) @@ -40,21 +15,24 @@ class FaceDetector(): self.targetimageSize = [224, 224] self.image = None - def detect_faces(self, frame): + def detect_faces(self, frame, model): self.resizeImage(frame) - outputs = self.model.run(self.image) + outputs, _ = model.faceDetection.run([], {'input': self.image}) + outputs = outputs[0] faceLocation = np.argmax(outputs[0].flatten()) x = faceLocation % 56 y = faceLocation // 56 + if outputs[0, y, x] < self.detection_threshold: - return [] + return None r = outputs[1, y, x] * 112. results= (((x * 4) - r, (y * 4) - r, r*2,r*2)) results = np.array(results).astype(np.float32) - results[ [0,2]] *= frame.shape[1] / 224. - results[ [1,3]] *= frame.shape[0] / 224. + results[[0,2]] *= frame.shape[1] / 224. + results[[1,3]] *= frame.shape[0] / 224. + return results def resizeImage(self, frame): diff --git a/landmarks.py b/landmarks.py index 8bb0f69..b6ce795 100644 --- a/landmarks.py +++ b/landmarks.py @@ -28,37 +28,13 @@ def matrix_to_quaternion(m): #not because I didn't try to improve things #but because I caused problems when I did class Landmarks(): - def __init__(self, width, height, model_type, threshold): - #removed all the tracking levels below 0 - #because they were so feature bare that they just didn't seem usable - #I beleive the selection should be equivalent to VTS, just 0 indexed - self.model_type = model_type - self.models = [ - "lm_model0_opt.onnx", - "lm_model1_opt.onnx", - "lm_model2_opt.onnx", - "lm_model3_opt.onnx", - "lm_model4_opt.onnx"] - - model = self.models[self.model_type] - model_base_path = os.path.join(os.path.dirname(__file__), os.path.join("models")) - providersList = onnxruntime.capi._pybind_state.get_available_providers() - options = onnxruntime.SessionOptions() - options.inter_op_num_threads = 1 - #the only num_threads that seems to matter - #Doesn't seem to scale well beyond 2 - options.intra_op_num_threads = 1 - options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL - options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL - options.log_severity_level = 3 - - self.session = onnxruntime.InferenceSession(os.path.join(model_base_path, model), sess_options=options, providers=providersList) - self.input_name = self.session.get_inputs()[0].name + def __init__(self, width, height, threshold): self.camera = np.array([[width, 0, width/2], [0, width, height/2], [0, 0, 1]], np.float32) self.inverse_camera = np.linalg.inv(self.camera) self.width = width self.height = height + self.confidence = 0 self.face_3d = np.array([ [ 0.4551769692672 , 0.300895790030204, -0.764429433974752], @@ -138,10 +114,11 @@ class Landmarks(): self.fail_count = 0 - def run(self, crop, crop_info): - output = self.session.run([], {self.input_name: crop})[0][0] - conf, lms = self.landmarks(output, crop_info) - return (conf, lms) + def run(self, models, crop, crop_info): + + output = models.landmarks.run([], {"input": crop})[0][0] + self.confidence, lms = self.landmarks(output, crop_info) + return ( lms) def landmarks(self, tensor, crop_info): crop_x1, crop_y1, scale_x, scale_y = crop_info diff --git a/tracker.py b/tracker.py index 219bdff..da50a76 100644 --- a/tracker.py +++ b/tracker.py @@ -1,12 +1,12 @@ +import os import numpy as np +import onnxruntime import cv2 import time import facedetection import eyes import landmarks import face -import featureExtractor -import emilianaFeatureExtractor #this file is so much smaller than it used to be, lol @@ -21,38 +21,63 @@ def clamp_to_im(pt, w, h): #8 times per frame, but that only accounts for 0.005m y = h-1 return (int(x),int(y+1)) + + +class Models(): + def __init__(self, model_type=3): + self.model_type = model_type + self.models = [ + "lm_model0_opt.onnx", + "lm_model1_opt.onnx", + "lm_model2_opt.onnx", + "lm_model3_opt.onnx", + "lm_model4_opt.onnx"] + + model = self.models[self.model_type] + model_base_path = os.path.join(os.path.dirname(__file__), os.path.join("models")) + providersList = onnxruntime.capi._pybind_state.get_available_providers() + options = onnxruntime.SessionOptions() + options.inter_op_num_threads = 1 + options.intra_op_num_threads = 1 + options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL + options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + options.log_severity_level = 3 + + self.landmarks = onnxruntime.InferenceSession(os.path.join(model_base_path, model), sess_options=options, providers=providersList) + + self.faceDetection = onnxruntime.InferenceSession(os.path.join(model_base_path, "mnv3_detection_opt.onnx"), sess_options=options, providers=providersList) + + self.gazeTracker = onnxruntime.InferenceSession(os.path.join(model_base_path, "mnv3_gaze32_split_opt.onnx"), sess_options=options, providers=providersList) + class Tracker(): def __init__(self, width, height, messageQueue, featureType, model_type=3, detection_threshold=0.6, threshold=0.6, silent=False): self.FaceDetector = facedetection.FaceDetector(detection_threshold = detection_threshold) self.EyeTracker = eyes.EyeTracker() - self.Landmarks = landmarks.Landmarks(width, height, model_type, threshold) - if featureType == 0: - self.FeatureExtractor = featureExtractor.FeatureExtractor() - else: - self.FeatureExtractor = emilianaFeatureExtractor.FeatureExtractor() - + self.Landmarks = landmarks.Landmarks(width, height, threshold) + self.model = Models(model_type = model_type) # Image normalization constants self.mean = np.float32(np.array([-2.1179, -2.0357, -1.8044])) self.std = np.float32(np.array([0.0171, 0.0175, 0.0174])) - + self.targetimageSize = [224, 224] self.width = width self.height = height self.threshold = threshold self.face = None - self.face_info = [face.FaceInfo(0, self, featureType)] + self.face_info = face.FaceInfo(featureType) self.messageQueue = messageQueue + #scales cropped frame before sending to landmarks def preprocess(self, im): - im = cv2.resize(im, (224, 224), interpolation=cv2.INTER_CUBIC) * self.std + self.mean + im = cv2.resize(im ,self.targetimageSize , interpolation=cv2.INTER_CUBIC) * self.std + self.mean im = np.expand_dims(im, 0) return np.transpose(im, (0,3,1,2)) - def cropFace(self,face, im): + def cropFace(self, im): duration_pp = 0.0 - x,y,w,h = face + x,y,w,h = self.face wint = int(w * 0.1) hint = int(h * 0.125) @@ -94,30 +119,29 @@ class Tracker(): if self.face is None: start_fd = time.perf_counter() - self.face = self.FaceDetector.detect_faces(frame) + self.face = self.FaceDetector.detect_faces(frame, self.model) duration_fd = 1000 * (time.perf_counter() - start_fd) if self.face is None: return self.early_exit("No faces found", start) - crop, crop_info, duration_pp = self.cropFace(self.face, frame) + crop, crop_info, duration_pp = self.cropFace(frame) #Early exit if crop fails, If the crop fails there's nothing to track if crop is None: return self.early_exit("No valid crops", start) - start_model = time.perf_counter() - conf, lms = self.Landmarks.run(crop , crop_info ) + lms = self.Landmarks.run(self.model, crop, crop_info) #Early exit if below confidence threshold - if conf < self.threshold: + if self.Landmarks.confidence < self.threshold: return self.early_exit("Confidence below threshold", start) - eye_state = self.EyeTracker.get_eye_state(frame, lms) + eye_state = self.EyeTracker.get_eye_state(self.model, frame, lms) - self.face_info[0].update((conf, (lms, eye_state)), np.array(lms)[:, 0:2].mean(0)) + self.face_info.update((self.Landmarks.confidence, (lms, eye_state)), np.array(lms)[:, 0:2].mean(0)) duration_model = 1000 * (time.perf_counter() - start_model) start_pnp = time.perf_counter() - face_info = self.face_info[0] + face_info = self.face_info if face_info.alive: face_info.success, face_info.quaternion, face_info.euler, face_info.pnp_error, face_info.pts_3d, face_info.lms = self.Landmarks.estimate_depth(face_info) diff --git a/vts.py b/vts.py index b041ac2..aa0442f 100644 --- a/vts.py +++ b/vts.py @@ -30,7 +30,7 @@ class VTS(): print(f"Confidence: {face.conf:.4f} / 3D fitting error: {face.pnp_error:.4f} / Eyes: {left_state}, {right_state}") now = time.time() packet.extend(bytearray(struct.pack("d", now))) - packet.extend(bytearray(struct.pack("i", face.id))) + packet.extend(bytearray(struct.pack("i", 0))) packet.extend(bytearray(struct.pack("f", self.width))) packet.extend(bytearray(struct.pack("f", self.height))) packet.extend(bytearray(struct.pack("f", face.eye_blink[0])))