From 994d3fc880bcb11de00c03622165f1fb58a97099 Mon Sep 17 00:00:00 2001 From: Cowwan Date: Mon, 5 Feb 2024 16:54:50 -0500 Subject: [PATCH] put in the new code --- eyes.py | 156 ++++++ face.py | 143 ++++++ facedetection.py | 52 ++ facetracker.py | 593 +++++++--------------- feature.py | 108 ++++ landmarks.py | 266 ++++++++++ preview.py | 21 + tracker.py | 1260 ++++------------------------------------------ vts.py | 72 +++ webcam.py | 104 ++++ 10 files changed, 1196 insertions(+), 1579 deletions(-) create mode 100644 eyes.py create mode 100644 face.py create mode 100644 facedetection.py create mode 100644 feature.py create mode 100644 landmarks.py create mode 100644 preview.py create mode 100644 vts.py create mode 100644 webcam.py diff --git a/eyes.py b/eyes.py new file mode 100644 index 0000000..eee19d0 --- /dev/null +++ b/eyes.py @@ -0,0 +1,156 @@ +import os +import numpy as np +import math +import cv2 +import onnxruntime +import time + + + +def clamp_to_im(pt, w, h): #8 times per frame, but that only accounts for 0.005ms + x=max(pt[0],0) + y=max(pt[1],0) + if x >= w: + x = w-1 + if y >= h: + y = h-1 + return (int(x),int(y+1)) + +def rotate(origin, point, a): #runs 22 times per frame + x, y = point - origin + cosa = math.cos(-a) + sina = math.sin(-a) + qx = origin[0] + cosa * x - sina * y + qy = origin[1] + sina * x + cosa * y + return qx, qy + +def clamp (value, minimum, maxium): + return max(min(value,maxium),minimum) + +def rotate_image(image, a, center): #twice per frame, 0.2ms - 0.25ms each, improved to 0.15ms - 0.25ms + M = cv2.getRotationMatrix2D(center, math.degrees(a), 1.0) + (h, w) = image.shape[:2] + image = cv2.warpAffine(image, M, (w, h), flags = cv2.INTER_CUBIC) + return image + +class EyeTracker(): + def __init__(self): + self.last_eye_state = [[0.,0.],[0.,0.]] + self.average_eye_confidence = [1.0,1.0] + self.average_eye_confidence_variance = [0.0,0.0] + self.eye_std_dev = [0.0,0.0] + self.eye_tracking_frames = 0 + self.mean = np.float32(np.array([-2.1179, -2.0357, -1.8044])) + self.std = np.float32(np.array([0.0171, 0.0175, 0.0174])) + + 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,face_frame, lms, offset): + eye_state = [[1.,0.,0.,0.],[1.,0.,0.,0.]] #Array constructed early so it can be used for an early exit + self.eye_tracking_frames+=1 + + #I've tried to thread this because they don't need to happen in sequence + #but it doesn't do any good because python threads kinda suck + #I'd need a persistent right eye and left eye process and that feels excessive + right_eye, REyeInfo = self.prepare_eye(face_frame, lms[36,0:2], lms[39,0:2], False) + left_eye, LEyeInfo = self.prepare_eye(face_frame,lms[42,0:2], lms[45,0:2], True) + + if right_eye is None or left_eye is None: + return eye_state #Early exit if one of the eyes doesn't have data + both_eyes = np.concatenate((right_eye, left_eye)) + + results = self.gaze_model.run([], {"input": both_eyes})[0] + eye_state[0] = self.calculateEye(results[0], REyeInfo, offset,0) + eye_state[1] = self.calculateEye(results[1], LEyeInfo, offset,1) + return eye_state + + def calculateStandardDeviation(self, conf, eye): + + avgRatio = 1/self.eye_tracking_frames + self.average_eye_confidence[eye] = (self.average_eye_confidence[eye]*(1-avgRatio))+ (conf * avgRatio) + + avgRatio = 1/(self.eye_tracking_frames + 1) + self.average_eye_confidence_variance[eye] = pow(((conf - self.average_eye_confidence[eye]) * avgRatio) + (self.average_eye_confidence_variance[eye] * (1-avgRatio)), 2) + + self.eye_std_dev[eye] = math.sqrt(self.average_eye_confidence_variance[eye]) + + def calculateEye(self, results, eye_info , offset, eye): + + e_x, e_y, scale, reference, angles = eye_info + confidenceThreshold = self.average_eye_confidence[eye] - 2 * self.eye_std_dev[eye] + + m = results[0].argmax() + x = m // 8 + y = m % 8 + + p=results[1][x, y] + p = clamp(p, 0.00001, 0.9999) + off_x = math.log(p/(1-p)) + eye_x = 4.0 *(x + off_x) + + p=results[2][x, y] + p = clamp(p, 0.00001, 0.9999) + off_y = math.log(p/(1-p)) + eye_y = 4.0 * (y + off_y) + + #I'm proud of this + #it's my "anti-Cali" functionality + #if eye movements are below 2 standard deviations of the average the movement is severely limited + #the eyes are handled independenly because my testing showed that one eye tended to have a higher condifence than the other + limit = max((results[0][x,y]- 0.5)*(2/confidenceThreshold),0) + if results[0][x,y] < confidenceThreshold: + Delta = eye_y - self.last_eye_state[eye][0] + eye_y = self.last_eye_state[eye][0] + clamp(Delta, -limit, limit) + Delta = eye_x - self.last_eye_state[eye][1] + eye_x = self.last_eye_state[eye][1] + clamp(Delta, -limit, limit) + + self.last_eye_state[eye] = [eye_y, eye_x] + + if eye == 1: + eye_x = (32. - eye_x) + eye_x = e_x + scale[0] * eye_x + eye_y = e_y + scale[1] * eye_y + + eye_x, eye_y = rotate(reference, (eye_x, eye_y), -angles) + eye_x, eye_y = (eye_x, eye_y) + offset + self.calculateStandardDeviation(results[0][x,y], eye) + return [1.0, eye_y, eye_x, results[0][x,y]] + + def prepare_eye(self,im, outer_pt, inner_pt, flip): + (x1, y1), (x2, y2), center, reference, a = self.corners_to_eye((outer_pt, inner_pt), im.shape[1], im.shape[0]) + #rotating an image is expensive and reduces clarity + #so I just don't if it's a relatively small angle + if math.degrees(a) > 7.5 and math.degrees(a) < 352.5: + im = rotate_image(im, a, reference) + + im = im[int(y1):int(y2), int(x1):int(x2)] + if np.prod(im.shape) < 1: + return None, None, None, None, None, None + if flip: + im = cv2.flip(im, 1) + scale = [(x2 - x1)/ 32., (y2 - y1)/ 32.] + im = cv2.resize(im, (32, 32), interpolation=cv2.INTER_CUBIC) + im = im.astype(np.float32) * self.std + self.mean + im = np.expand_dims(im, 0) + im = np.transpose(im, (0,3,2,1)) + return [im, [x1, y1, scale, reference, a]] + + def corners_to_eye(self, corners, w, h): + c1, c2 = np.array(corners, dtype=float) + a = math.atan2(*(c2 - c1)[::-1]) % (2*math.pi) + c2 = rotate(c1, c2, a) + center = (c1 + c2) / 2.0 + radius = np.linalg.norm(c1 - c2) + radius = [radius * 0.7, radius * 0.6] + upper_left = clamp_to_im(center - radius, w, h) + lower_right = clamp_to_im(center + radius, w, h) + return upper_left, lower_right, center, c1, a diff --git a/face.py b/face.py new file mode 100644 index 0000000..81e0e7d --- /dev/null +++ b/face.py @@ -0,0 +1,143 @@ +import numpy as np +import math +import feature + +class FaceInfo(): + def __init__(self, id, tracker): + self.id = id + self.tracker = tracker + self.face_3d = np.array([ + [ 0.4551769692672 , 0.300895790030204, -0.764429433974752], + [ 0.448998827123556, 0.166995837790733, -0.765143004071253], + [ 0.437431554952677, 0.022655479179981, -0.739267175112735], + [ 0.415033422928434, -0.088941454648772, -0.747947437846473], + [ 0.389123587370091, -0.232380029794684, -0.704788385327458], + [ 0.334630113904382, -0.361265387599081, -0.615587579236862], + [ 0.263725112132858, -0.460009725616771, -0.491479221041573], + [ 0.16241621322721 , -0.558037146073869, -0.339445180872282], + [ 0. , -0.621079019321682, -0.287294770748887], + [-0.16241621322721 , -0.558037146073869, -0.339445180872282], + [-0.263725112132858, -0.460009725616771, -0.491479221041573], + [-0.334630113904382, -0.361265387599081, -0.615587579236862], + [-0.389123587370091, -0.232380029794684, -0.704788385327458], + [-0.415033422928434, -0.088941454648772, -0.747947437846473], + [-0.437431554952677, 0.022655479179981, -0.739267175112735], + [-0.448998827123556, 0.166995837790733, -0.765143004071253], + [-0.4551769692672 , 0.300895790030204, -0.764429433974752], + [ 0.385529968662985, 0.402800553948697, -0.310031082540741], + [ 0.322196658344302, 0.464439136821772, -0.250558059367669], + [ 0.25409760441282 , 0.46420381416882 , -0.208177722146526], + [ 0.186875436782135, 0.44706071961879 , -0.145299823706503], + [ 0.120880983543622, 0.423566314072968, -0.110757158774771], + [-0.120880983543622, 0.423566314072968, -0.110757158774771], + [-0.186875436782135, 0.44706071961879 , -0.145299823706503], + [-0.25409760441282 , 0.46420381416882 , -0.208177722146526], + [-0.322196658344302, 0.464439136821772, -0.250558059367669], + [-0.385529968662985, 0.402800553948697, -0.310031082540741], + [ 0. , 0.293332603215811, -0.137582088779393], + [ 0. , 0.194828701837823, -0.069158109325951], + [ 0. , 0.103844017393155, -0.009151819844964], + [ 0. , 0. , 0. ], + [ 0.080626352317973, -0.041276068128093, -0.134161035564826], + [ 0.046439347377934, -0.057675223874769, -0.102990627164664], + [ 0. , -0.068753126205604, -0.090545348482397], + [-0.046439347377934, -0.057675223874769, -0.102990627164664], + [-0.080626352317973, -0.041276068128093, -0.134161035564826], + [ 0.315905195966084, 0.298337502555443, -0.285107407636464], + [ 0.275252345439353, 0.312721904921771, -0.244558251170671], + [ 0.176394511553111, 0.311907184376107, -0.219205360345231], + [ 0.131229723798772, 0.284447361805627, -0.234239149487417], + [ 0.184124948330084, 0.260179585304867, -0.226590776513707], + [ 0.279433549294448, 0.267363071770222, -0.248441437111633], + [-0.131229723798772, 0.284447361805627, -0.234239149487417], + [-0.176394511553111, 0.311907184376107, -0.219205360345231], + [-0.275252345439353, 0.312721904921771, -0.244558251170671], + [-0.315905195966084, 0.298337502555443, -0.285107407636464], + [-0.279433549294448, 0.267363071770222, -0.248441437111633], + [-0.184124948330084, 0.260179585304867, -0.226590776513707], + [ 0.121155252430729, -0.208988660580347, -0.160606287940521], + [ 0.041356305910044, -0.194484199722098, -0.096159882202821], + [ 0. , -0.205180167345702, -0.083299217789729], + [-0.041356305910044, -0.194484199722098, -0.096159882202821], + [-0.121155252430729, -0.208988660580347, -0.160606287940521], + [-0.132325402795928, -0.290857984604968, -0.187067868218105], + [-0.064137791831655, -0.325377847425684, -0.158924039726607], + [ 0. , -0.343742581679188, -0.113925986025684], + [ 0.064137791831655, -0.325377847425684, -0.158924039726607], + [ 0.132325402795928, -0.290857984604968, -0.187067868218105], + [ 0.181481567104525, -0.243239316141725, -0.231284988892766], + [ 0.083999507750469, -0.239717753728704, -0.155256465640701], + [ 0. , -0.256058040176369, -0.0950619498899 ], + [-0.083999507750469, -0.239717753728704, -0.155256465640701], + [-0.181481567104525, -0.243239316141725, -0.231284988892766], + [-0.074036069749345, -0.250689938345682, -0.177346470406188], + [ 0. , -0.264945854681568, -0.112349967428413], + [ 0.074036069749345, -0.250689938345682, -0.177346470406188], + # Pupils and eyeball centers + [ 0.257990002632141, 0.276080012321472, -0.219998998939991], + [-0.257990002632141, 0.276080012321472, -0.219998998939991], + [ 0.257990002632141, 0.276080012321472, -0.324570998549461], + [-0.257990002632141, 0.276080012321472, -0.324570998549461] + ], np.float32) + + + self.reset() + self.alive = False + self.coord = None + self.base_scale_v = self.face_3d[27:30, 1] - self.face_3d[28:31, 1] + self.base_scale_h = np.abs(self.face_3d[[0, 36, 42], 0] - self.face_3d[[16, 39, 45], 0]) + + def reset(self): + self.alive = False + self.conf = None + self.lms = None + self.eye_state = None + self.rotation = np.array([0.0, 0.0, 0.0], np.float32) + self.translation = np.array([0.0, 0.0, 0.0], np.float32) + self.success = None + self.quaternion = None + self.euler = None + self.pts_3d = None + self.eye_blink = None + self.pnp_error = 0 + self.features = feature.FeatureExtractor() + self.current_features = {} + self.contour = np.zeros((21,3)) + self.update_contour() + + def update(self, result, coord): + if result is None: + self.reset() + else: + self.conf, (self.lms, self.eye_state) = result + self.coord = coord + self.alive = True + + def update_contour(self): + self.contour = self.face_3d[[0,1,8,15,16,27,28,29,30,31,32,33,34,35]] + + def normalize_pts3d(self, pts_3d): + # Calculate angle using nose + pts_3d[:, 0:2] -= pts_3d[30, 0:2] + alpha = (math.atan2(*((pts_3d[27, 0:2]) - pts_3d[30, 0:2])[::-1]) % (2*math.pi)) + cosalpha = math.cos(alpha-(math.pi/2)) + sinalpha = math.sin(alpha-(math.pi/2)) + + R = np.matrix([[cosalpha, -sinalpha], [sinalpha, cosalpha]]) + pts_3d[:, 0:2] = (pts_3d - pts_3d[30])[:, 0:2].dot(R) + pts_3d[30, 0:2] + + # Vertical scale + pts_3d[:, 1] /= np.mean((pts_3d[27:30, 1] - pts_3d[28:31, 1]) / self.base_scale_v) + + # Horizontal scale + pts_3d[:, 0] /= np.mean(np.abs(pts_3d[[0, 36, 42], 0] - pts_3d[[16, 39, 45], 0]) / self.base_scale_h) + + return pts_3d + + def adjust_3d(self): + if self.conf < 0.4 or self.pnp_error > 300: + return + + self.pts_3d = self.normalize_pts3d(self.pts_3d) + self.current_features = self.features.update(self.pts_3d[:, 0:2]) + self.eye_blink = [self.current_features["eye_r"],self.current_features["eye_l"] ] diff --git a/facedetection.py b/facedetection.py new file mode 100644 index 0000000..c4ab90d --- /dev/null +++ b/facedetection.py @@ -0,0 +1,52 @@ +import os +import numpy as np +import cv2 +import onnxruntime + + +#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 FaceDetector(): + def __init__(self, detection_threshold=0.6): + + 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) + + 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.detection_threshold = detection_threshold + + def detect_faces(self, frame): + + im = cv2.resize(frame, (224, 224), interpolation=cv2.INTER_LINEAR) * self.std + self.mean + im = np.expand_dims(im, 0) + im = np.transpose(im, (0,3,1,2)) + + outputs, maxpool = self.detection.run([], {'input': im}) + + outputs[0, 0, outputs[0, 0] != maxpool[0, 0]] = 0 + detections = np.flip(np.argsort(outputs[0,0].flatten())) + results = [] + if len(detections) == 0: + return results + + x = detections[0] % 56 + y = detections[0] // 56 + if outputs[0, 0, y, x] < self.detection_threshold: + return results + + r = outputs[0, 1, y, x] * 112. + results.append(((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. + return results diff --git a/facetracker.py b/facetracker.py index a30628e..8d322e3 100644 --- a/facetracker.py +++ b/facetracker.py @@ -1,455 +1,206 @@ -import copy +#import copy import os -import sys +#limits the threads used by numpy to 1 +#setting this and the landmarks model to use 1 means the cpu usage is pretty low +os.environ["OMP_NUM_THREADS"] = str(1) import argparse import traceback -import gc +import threading +import queue +import time +from tracker import Tracker +import multiprocessing +import webcam +import vts +import preview + +#overall, my changes seem to have memory usage up from about 140mb to 180mb +#but cpu usage seems to be down +#and it's more resilent to low performance scenarios parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("-i", "--ip", help="Set IP address for sending tracking data", default="127.0.0.1") parser.add_argument("-p", "--port", type=int, help="Set port for sending tracking data", default=11573) -if os.name == 'nt': - parser.add_argument("-l", "--list-cameras", type=int, help="Set this to 1 to list the available cameras and quit, set this to 2 or higher to output only the names", default=0) - parser.add_argument("-a", "--list-dcaps", type=int, help="Set this to -1 to list all cameras and their available capabilities, set this to a camera id to list that camera's capabilities", default=None) - parser.add_argument("-W", "--width", type=int, help="Set camera and raw RGB width", default=640) - parser.add_argument("-H", "--height", type=int, help="Set camera and raw RGB height", default=360) - parser.add_argument("-D", "--dcap", type=int, help="Set which device capability line to use or -1 to use the default camera settings (FPS still need to be set separately)", default=None) - parser.add_argument("-B", "--blackmagic", type=int, help="When set to 1, special support for Blackmagic devices is enabled", default=0) -else: - parser.add_argument("-W", "--width", type=int, help="Set raw RGB width", default=640) - parser.add_argument("-H", "--height", type=int, help="Set raw RGB height", default=360) +parser.add_argument("-W", "--width", type=int, help="Set raw RGB width", default=640) +parser.add_argument("-H", "--height", type=int, help="Set raw RGB height", default=360) parser.add_argument("-F", "--fps", type=int, help="Set camera frames per second", default=24) -parser.add_argument("-c", "--capture", help="Set camera ID (0, 1...) or video file", default="0") +parser.add_argument("-c", "--capture", help="Set camera ID (0, 1...) or video file", default=0) parser.add_argument("-M", "--mirror-input", action="store_true", help="Process a mirror image of the input video") -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 tracking", default=None) +parser.add_argument("-t", "--threshold", type=float, help="Set minimum confidence threshold for face tracking", default=0.6) parser.add_argument("-d", "--detection-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) -parser.add_argument("--scan-retinaface", type=int, help="When set to 1, scanning for additional faces will be performed using RetinaFace in a background thread, otherwise a simpler, faster face detection mechanism is used. When the maximum number of faces is 1, this option does nothing.", default=0) -parser.add_argument("--scan-every", type=int, help="Set after how many frames a scan for new faces should run", default=3) -parser.add_argument("--discard-after", type=int, help="Set the how long the tracker should keep looking for lost faces", default=10) -parser.add_argument("--max-feature-updates", type=int, help="This is the number of seconds after which feature min/max/medium values will no longer be updated once a face has been detected.", default=900) -parser.add_argument("--no-3d-adapt", type=int, help="When set to 1, the 3D face model will not be adapted to increase the fit", default=1) -parser.add_argument("--try-hard", type=int, help="When set to 1, the tracker will try harder to find a face", default=0) -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("--video-scale", type=int, help="This is a resolution scale factor applied to the saved AVI file", default=1, choices=[1,2,3,4]) -parser.add_argument("--video-fps", type=float, help="This sets the frame rate of the output AVI file", default=24) -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, except for model 4, which is wink optimized. Models 1 and 0 tend to be too rigid for expression and blink detection. Model -2 is roughly equivalent to model 1, but faster. Model -3 is between models 0 and -1.", default=3, choices=[-3, -2, -1, 0, 1, 2, 3, 4]) -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, gaze tracking is 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) -parser.add_argument("--repeat-video", type=int, help="When set to 1 and a video file was specified with -c, the tracker will loop the video until interrupted", default=0) -parser.add_argument("--dump-points", type=str, help="When set to a filename, the current face 3D points are made symmetric and dumped to the given file when quitting the visualization with the \"q\" key", default="") -parser.add_argument("--benchmark", type=int, help="When set to 1, the different tracking models are benchmarked, starting with the best and ending with the fastest and with gaze tracking disabled for models with negative IDs", default=0) -if os.name == 'nt': - parser.add_argument("--use-dshowcapture", type=int, help="When set to 1, libdshowcapture will be used for video input instead of OpenCV", default=1) - parser.add_argument("--blackmagic-options", type=str, help="When set, this additional option string is passed to the blackmagic capture library", default=None) - parser.add_argument("--priority", type=int, help="When set, the process priority will be changed", default=None, choices=[0, 1, 2, 3, 4, 5]) +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, except for model 4, which is wink optimized.", default=3, choices=[0, 1, 2, 3, 4]) +parser.add_argument("--preview", type=int, help="Preview the frames sent to the tracker", default=0) args = parser.parse_args() -os.environ["OMP_NUM_THREADS"] = str(args.max_threads) -class OutputLog(object): - def __init__(self, fh, output): - self.fh = fh - self.output = output - def write(self, buf): - if self.fh is not None: - self.fh.write(buf) - self.output.write(buf) - self.flush() - def flush(self): - if self.fh is not None: - self.fh.flush() - self.output.flush() -output_logfile = None -if args.log_output != "": - output_logfile = open(args.log_output, "w") -sys.stdout = OutputLog(output_logfile, sys.stdout) -sys.stderr = OutputLog(output_logfile, sys.stderr) - -if os.name == 'nt': - import dshowcapture - if args.blackmagic == 1: - dshowcapture.set_bm_enabled(True) - if args.blackmagic_options is not None: - dshowcapture.set_options(args.blackmagic_options) - if args.priority is not None: - import psutil - classes = [psutil.IDLE_PRIORITY_CLASS, psutil.BELOW_NORMAL_PRIORITY_CLASS, psutil.NORMAL_PRIORITY_CLASS, psutil.ABOVE_NORMAL_PRIORITY_CLASS, psutil.HIGH_PRIORITY_CLASS, psutil.REALTIME_PRIORITY_CLASS] - p = psutil.Process(os.getpid()) - p.nice(classes[args.priority]) - -if os.name == 'nt' and (args.list_cameras > 0 or args.list_dcaps is not None): - cap = dshowcapture.DShowCapture() - info = cap.get_info() - unit = 10000000.; - if args.list_dcaps is not None: - formats = {0: "Any", 1: "Unknown", 100: "ARGB", 101: "XRGB", 200: "I420", 201: "NV12", 202: "YV12", 203: "Y800", 300: "YVYU", 301: "YUY2", 302: "UYVY", 303: "HDYC (Unsupported)", 400: "MJPEG", 401: "H264" } - for cam in info: - if args.list_dcaps == -1: - type = "" - if cam['type'] == "Blackmagic": - type = "Blackmagic: " - print(f"{cam['index']}: {type}{cam['name']}") - if args.list_dcaps != -1 and args.list_dcaps != cam['index']: - continue - for caps in cam['caps']: - format = caps['format'] - if caps['format'] in formats: - format = formats[caps['format']] - if caps['minCX'] == caps['maxCX'] and caps['minCY'] == caps['maxCY']: - print(f" {caps['id']}: Resolution: {caps['minCX']}x{caps['minCY']} FPS: {unit/caps['maxInterval']:.3f}-{unit/caps['minInterval']:.3f} Format: {format}") - else: - print(f" {caps['id']}: Resolution: {caps['minCX']}x{caps['minCY']}-{caps['maxCX']}x{caps['maxCY']} FPS: {unit/caps['maxInterval']:.3f}-{unit/caps['minInterval']:.3f} Format: {format}") - else: - if args.list_cameras == 1: - print("Available cameras:") - for cam in info: - type = "" - if cam['type'] == "Blackmagic": - type = "Blackmagic: " - if args.list_cameras == 1: - print(f"{cam['index']}: {type}{cam['name']}") - else: - print(f"{type}{cam['name']}") - cap.destroy_capture() - sys.exit(0) - -import numpy as np -import time -import cv2 -import socket -import struct -import json -from input_reader import InputReader, VideoReader, DShowCaptureReader, try_int -from tracker import Tracker, get_model_base_path - -if args.benchmark > 0: - model_base_path = get_model_base_path(args.model_dir) - im = cv2.imread(os.path.join(model_base_path, "benchmark.bin"), cv2.IMREAD_COLOR) - results = [] - for model_type in [3, 2, 1, 0, -1, -2, -3]: - tracker = Tracker(224, 224, threshold=0.1, max_threads=args.max_threads, max_faces=1, discard_after=0, scan_every=0, silent=True, model_type=model_type, model_dir=args.model_dir, no_gaze=(model_type == -1), detection_threshold=0.1, use_retinaface=0, max_feature_updates=900, static_model=True if args.no_3d_adapt == 1 else False) - tracker.detected = 1 - tracker.faces = [(0, 0, 224, 224)] - total = 0.0 - for i in range(100): - start = time.perf_counter() - r = tracker.predict(im) - total += time.perf_counter() - start - print(1. / (total / 100.)) - sys.exit(0) +def messaging(messageQueue): + while True: + message = messageQueue.get() + print(message) +#processing args +#For Vtube Studio: ip 127.0.0.1 port 11573 target_ip = args.ip target_port = args.port - -if args.faces >= 40: - print("Transmission of tracking data over network is not supported with 40 or more faces.") - fps = args.fps -dcap = None -use_dshowcapture_flag = False -if os.name == 'nt': - dcap = args.dcap - use_dshowcapture_flag = True if args.use_dshowcapture == 1 else False - input_reader = InputReader(args.capture, args.raw_rgb, args.width, args.height, fps, use_dshowcapture=use_dshowcapture_flag, dcap=dcap) - if args.dcap == -1 and type(input_reader) == DShowCaptureReader: - fps = min(fps, input_reader.device.get_fps()) -else: - input_reader = InputReader(args.capture, args.raw_rgb, args.width, args.height, fps) -if type(input_reader.reader) == VideoReader: - fps = 0 +height = args.height +width = args.width +previewFlag = (args.preview == 1) +mirror_input = args.mirror_input + +#I make a minor tweak to the framerate to help the camera out +#it's an inpercetible amount of time per frame +#but my hope is it keeps things from riding up against what the webcam can put out +target_duration = 1. / (fps-0.01) +targetBrightness = 0.55 +silent = (args.silent == 1) + + +#increating the webcam queue to compensate for longer frame times at different settings +frameQueueSize = 1 +if height > 480: + frameQueueSize = 3 + #My testing has shown that it can take upwards of 80ms when capturing images above 480p + #each frame is 33ms, so I'd need 3 frames, or 99ms to cover that without skipping a frame +if fps > 30: + frameQueueSize = frameQueueSize * 2 + #I'm just going to assume the webcam does not work well above 30fps + +#creating queues +frameQueue = multiprocessing.Queue(maxsize=frameQueueSize) +faceQueue = multiprocessing.Queue(maxsize=1) +messageQueue = queue.Queue() +packetQueue = queue.Queue() + +#creating threads +#probably doesn't need it's own thread, but it can't hurt +messageThread = threading.Thread(target = messaging, args = [ messageQueue],) +messageThread.daemon = True +messageThread.start() + +#I want to decouple the requests so they don't block anything else if they're slow +VTS = vts.VTS(target_ip, target_port, silent, height, width, packetQueue, messageQueue) +packetSenderThread = threading.Thread(target = VTS.packetSender, args = [],) +packetSenderThread.daemon = True +packetSenderThread.start() + +#creating processes +#not a very active process, it's blocked by the main thread 90% of the time +#but putting it in it's own process means there's a little bit less occupying the same core as the main thread +#which helps when the system is under load +webcamProcess = multiprocessing.Process(target=webcam.startProcess,args = (frameQueue, faceQueue, fps, targetBrightness, width, height, mirror_input,messageQueue, )) +webcamProcess.daemon = True +webcamProcess.start() + +#This is it's own process because it's super weird performance wise +if previewFlag: + previewFrameQueue = multiprocessing.Queue(maxsize=1) + previewProcess = multiprocessing.Process(target=preview.startProcess, args = (previewFrameQueue, target_duration,)) + previewProcess.daemon = True + previewProcess.start() -log = None -out = None -first = True -height = 0 -width = 0 -tracker = None -sock = None -total_tracking_time = 0.0 -tracking_time = 0.0 -tracking_frames = 0 frame_count = 0 +peak_frame_time=0.0 +total_active_time = 0.0 +total_run_time = 0.0 +frame_start = 0.0 +peak_time_between = 0.0 +peak_camera_latency = 0.0 +total_camera_latency = 0.0 +failures = 0 -features = ["eye_l", "eye_r", "eyebrow_steepness_l", "eyebrow_updown_l", "eyebrow_quirk_l", "eyebrow_steepness_r", "eyebrow_updown_r", "eyebrow_quirk_r", "mouth_corner_updown_l", "mouth_corner_inout_l", "mouth_corner_updown_r", "mouth_corner_inout_r", "mouth_open", "mouth_wide"] - -if args.log_data != "": - log = open(args.log_data, "w") - log.write("Frame,Time,Width,Height,FPS,Face,FaceID,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): - log.write(f",Point3D[{i}].X,Point3D[{i}].Y,Point3D[{i}].Z") - for feature in features: - log.write(f",{feature}") - log.write("\r\n") - log.flush() - -is_camera = args.capture == str(try_int(args.capture)) +tracker = Tracker(width, height, messageQueue, threshold=args.threshold, silent=silent, model_type=args.model, detection_threshold=args.detection_threshold) +#don't start until the webcam is ready +while frameQueue.qsize() < 1: + time.sleep(0.33) try: - attempt = 0 - frame_time = time.perf_counter() - target_duration = 0 - if fps > 0: - target_duration = 1. / float(fps) - repeat = args.repeat_video != 0 and type(input_reader.reader) == VideoReader - need_reinit = 0 - failures = 0 - source_name = input_reader.name - while repeat or input_reader.is_open(): - if not input_reader.is_open() or need_reinit == 1: - input_reader = InputReader(args.capture, args.raw_rgb, args.width, args.height, fps, use_dshowcapture=use_dshowcapture_flag, dcap=dcap) - if input_reader.name != source_name: - print(f"Failed to reinitialize camera and got {input_reader.name} instead of {source_name}.") - sys.exit(1) - need_reinit = 2 - time.sleep(0.02) - continue - if not input_reader.is_ready(): - time.sleep(0.02) - continue - - ret, frame = input_reader.read() - if ret and args.mirror_input: - frame = cv2.flip(frame, 1) - if not ret: - if repeat: - if need_reinit == 0: - need_reinit = 1 - continue - elif is_camera: - attempt += 1 - if attempt > 30: - break - else: - time.sleep(0.02) - if attempt == 3: - need_reinit = 1 - continue - else: - break; - - attempt = 0 - need_reinit = 0 + while True: + #clearing these so there's no stale data to confuse the checks + packet = None + face = None + frame_start = time.perf_counter() frame_count += 1 - now = time.time() + #if the camera process doesn't have a frame for us we just skip everything + #in my testing the camera was consistently waiting on the main process + #so if it's *that far behind* it probably needs a frame or two to catch up + if frameQueue.qsize() > 0: + frame, camera_latency = frameQueue.get() - if first: - first = False - height, width, channels = frame.shape - 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, model_dir=args.model_dir, no_gaze=False if args.gaze_tracking != 0 and args.model != -1 else True, detection_threshold=args.detection_threshold, use_retinaface=args.scan_retinaface, max_feature_updates=args.max_feature_updates, static_model=True if args.no_3d_adapt == 1 else False, try_hard=args.try_hard == 1) - if args.video_out is not None: - out = cv2.VideoWriter(args.video_out, cv2.VideoWriter_fourcc('F','F','V','1'), args.video_fps, (width * args.video_scale, height * args.video_scale)) + #If I don't wait a few frames to start tracking I get wild peak frame times, like 500ms + if frame_count > 5: + peak_camera_latency = max(camera_latency, peak_camera_latency) + total_camera_latency+= camera_latency - try: - inference_start = time.perf_counter() - faces = tracker.predict(frame) - if len(faces) > 0: - inference_time = (time.perf_counter() - inference_start) - total_tracking_time += inference_time - tracking_time += inference_time / len(faces) - tracking_frames += 1 - packet = bytearray() - detected = False - for face_num, f in enumerate(faces): - f = copy.copy(f) - f.id += args.face_id_offset - if f.eye_blink is None: - f.eye_blink = [1, 1] - right_state = "O" if f.eye_blink[0] > 0.30 else "-" - left_state = "O" if f.eye_blink[1] > 0.30 else "-" - if args.silent == 0: - print(f"Confidence[{f.id}]: {f.conf:.4f} / 3D fitting error: {f.pnp_error:.4f} / Eyes: {left_state}, {right_state}") - detected = True - if not f.success: - pts_3d = np.zeros((70, 3), np.float32) - packet.extend(bytearray(struct.pack("d", now))) - packet.extend(bytearray(struct.pack("i", f.id))) - packet.extend(bytearray(struct.pack("f", width))) - packet.extend(bytearray(struct.pack("f", height))) - packet.extend(bytearray(struct.pack("f", f.eye_blink[0]))) - packet.extend(bytearray(struct.pack("f", f.eye_blink[1]))) - packet.extend(bytearray(struct.pack("B", 1 if f.success else 0))) - packet.extend(bytearray(struct.pack("f", f.pnp_error))) - packet.extend(bytearray(struct.pack("f", f.quaternion[0]))) - packet.extend(bytearray(struct.pack("f", f.quaternion[1]))) - packet.extend(bytearray(struct.pack("f", f.quaternion[2]))) - packet.extend(bytearray(struct.pack("f", f.quaternion[3]))) - packet.extend(bytearray(struct.pack("f", f.euler[0]))) - packet.extend(bytearray(struct.pack("f", f.euler[1]))) - packet.extend(bytearray(struct.pack("f", f.euler[2]))) - packet.extend(bytearray(struct.pack("f", f.translation[0]))) - packet.extend(bytearray(struct.pack("f", f.translation[1]))) - packet.extend(bytearray(struct.pack("f", f.translation[2]))) - if log is not None: - log.write(f"{frame_count},{now},{width},{height},{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,0,255)) - for pt_num, (x,y,c) in enumerate(f.lms): - packet.extend(bytearray(struct.pack("f", y))) - packet.extend(bytearray(struct.pack("f", x))) - if log is not None: - log.write(f",{y},{x},{c}") - if pt_num == 66 and (f.eye_blink[0] < 0.30 or c < 0.20): - continue - if pt_num == 67 and (f.eye_blink[1] < 0.30 or c < 0.20): - continue - x = int(x + 0.5) - y = int(y + 0.5) - if args.visualize != 0 or out is not None: - 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, 255, 0) - if pt_num >= 66: - color = (255, 255, 0) - if not (x < 0 or y < 0 or x >= height or y >= width): - cv2.circle(frame, (y, x), 1, color, -1) - if args.pnp_points != 0 and (args.visualize != 0 or out is not None) and f.rotation is not None: - if args.pnp_points > 1: - projected = cv2.projectPoints(f.face_3d[0:66], f.rotation, f.translation, tracker.camera, tracker.dist_coeffs) - else: - projected = cv2.projectPoints(f.contour, f.rotation, f.translation, tracker.camera, tracker.dist_coeffs) - for [(x,y)] in projected[0]: - x = int(x + 0.5) - y = int(y + 0.5) - if not (x < 0 or y < 0 or x >= height or y >= width): - frame[int(x), int(y)] = (0, 255, 255) - x += 1 - if not (x < 0 or y < 0 or x >= height or y >= width): - frame[int(x), int(y)] = (0, 255, 255) - y += 1 - if not (x < 0 or y < 0 or x >= height or y >= width): - frame[int(x), int(y)] = (0, 255, 255) - x -= 1 - if not (x < 0 or y < 0 or x >= height or y >= width): - frame[int(x), int(y)] = (0, 255, 255) - for (x,y,z) in f.pts_3d: - packet.extend(bytearray(struct.pack("f", x))) - packet.extend(bytearray(struct.pack("f", -y))) - packet.extend(bytearray(struct.pack("f", -z))) - if log is not None: - log.write(f",{x},{-y},{-z}") - if f.current_features is None: - f.current_features = {} - for feature in features: - if not feature in f.current_features: - f.current_features[feature] = 0 - packet.extend(bytearray(struct.pack("f", f.current_features[feature]))) - if log is not None: - log.write(f",{f.current_features[feature]}") - if log is not None: - log.write("\r\n") - log.flush() + if previewFlag: + #there's no reason to ever wait on this process or give it much of a queue + if previewFrameQueue.qsize() < 1: + previewFrameQueue.put(frame) - if detected and len(faces) < 40: - sock.sendto(packet, (target_ip, target_port)) - - if out is not None: - video_frame = frame - if args.video_scale != 1: - video_frame = cv2.resize(frame, (width * args.video_scale, height * args.video_scale), interpolation=cv2.INTER_NEAREST) - out.write(video_frame) - if args.video_scale != 1: - del video_frame - - if args.visualize != 0: - cv2.imshow('OpenSeeFace Visualization', frame) - if cv2.waitKey(1) & 0xFF == ord('q'): - if args.dump_points != "" and faces is not None and len(faces) > 0: - np.set_printoptions(threshold=sys.maxsize, precision=15) - pairs = [ - (0, 16), - (1, 15), - (2, 14), - (3, 13), - (4, 12), - (5, 11), - (6, 10), - (7, 9), - (17, 26), - (18, 25), - (19, 24), - (20, 23), - (21, 22), - (31, 35), - (32, 34), - (36, 45), - (37, 44), - (38, 43), - (39, 42), - (40, 47), - (41, 46), - (48, 52), - (49, 51), - (56, 54), - (57, 53), - (58, 62), - (59, 61), - (65, 63) - ] - points = copy.copy(faces[0].face_3d) - for a, b in pairs: - x = (points[a, 0] - points[b, 0]) / 2.0 - y = (points[a, 1] + points[b, 1]) / 2.0 - z = (points[a, 2] + points[b, 2]) / 2.0 - points[a, 0] = x - points[b, 0] = -x - points[[a, b], 1] = y - points[[a, b], 2] = z - points[[8, 27, 28, 29, 33, 50, 55, 60, 64], 0] = 0.0 - points[30, :] = 0.0 - with open(args.dump_points, "w") as fh: - fh.write(repr(points)) + try: + face, face_center, face_radius = tracker.predict(frame) + failures = 0 + except Exception as e: + if e.__class__ == KeyboardInterrupt: + if not silent: + messageQueue.put("Quitting") + break + traceback.print_exc() + failures += 1 + if failures > 30: break - failures = 0 - except Exception as e: - if e.__class__ == KeyboardInterrupt: - if args.silent == 0: - print("Quitting") - break - traceback.print_exc() - failures += 1 - if failures > 30: - break - collected = False - del frame + if face is not None: + packet = VTS.preparePacket(face) + #this can maybe cause weirdness in the gamma calculations + #but I don't want to let that build a queue + if faceQueue.qsize() < 1: + faceQueue.put([face_center, face_radius]) + else: + messageQueue.put("Camera behind, skipping frame") + + duration = time.perf_counter() - frame_start + + total_active_time += duration + peak_frame_time = max(peak_frame_time, duration) + if duration < target_duration: + time.sleep(target_duration - duration) + else: + #all caps because if this is consistently happening it's an issue + messageQueue.put("CANNOT MAINTAIN FRAMERATE!") + + peak_time_between = max(peak_time_between, time.perf_counter() -frame_start) + total_run_time += time.perf_counter() -frame_start + + #If we don't have something to send to Vtube Studio we don't + if(packet is not None): + packetQueue.put(packet) + else: + messageQueue.put("No data sent to VTS") - duration = time.perf_counter() - frame_time - while duration < target_duration: - if not collected: - gc.collect() - collected = True - duration = time.perf_counter() - frame_time - sleep_time = target_duration - duration - if sleep_time > 0: - time.sleep(sleep_time) - duration = time.perf_counter() - frame_time - frame_time = time.perf_counter() except KeyboardInterrupt: - if args.silent == 0: + if not silent: print("Quitting") -input_reader.close() -if out is not None: - out.release() -cv2.destroyAllWindows() +#printing statistics on close +#it makes identifying problems easier +#480p generally results in camera latencies between 0.5ms and 2.0ms +print(f"Peak camera latency: {(peak_camera_latency*1000):.3f}ms") +print(f"Average camera latency: {(total_camera_latency*1000/(frame_count-5)):.3f}ms") +#The longest time a full cycle of the main loop takes, including sleep +#indicates issues when higher than average time between frames by a significant amount +print(f"Peak time between frames: {(peak_time_between*1000):.3f}ms") +#at 30fps this should be a hair above 33.333ms +print(f"Average time between frames: {(total_run_time*1000/(frame_count)):.3f}ms") +#These measure how long the main loop takes *not including sleep* +#tends to be slightly longer than the average time given by the tracking status by 1ms to 2ms +#because it includes detection frames and all the other minor operations that need to happen every frame +print(f"Peak frame time: {(peak_frame_time*1000):.3f}ms") +print(f"Average frame time: { ((total_active_time)*1000/(frame_count)):.3f}ms") +print(f"Run time (seconds): {total_run_time:.2f} s\nFrames: {frame_count}") -if args.silent == 0 and tracking_frames > 0: - average_tracking_time = 1000 * tracking_time / tracking_frames - print(f"Average tracking time per detected face: {average_tracking_time:.2f} ms") - print(f"Tracking time: {total_tracking_time:.3f} s\nFrames: {tracking_frames}") +webcamProcess.terminate() +if previewFlag: + previewProcess.terminate() +os._exit(0) diff --git a/feature.py b/feature.py new file mode 100644 index 0000000..e648e74 --- /dev/null +++ b/feature.py @@ -0,0 +1,108 @@ +import os +import numpy as np +import math +import cv2 + +#got tired of trying to remember which order I had to do the nested min and max +def clamp (value, minimum, maxium): + return max(min(value,maxium),minimum) + +def rotate(origin, point, a): + x, y = point - origin + cosa = math.cos(-a) + sina = math.sin(-a) + qx = origin[0] + cosa * x - sina * y + qy = origin[1] + sina * x + cosa * y + return qx, qy + +#I redid a lot of this so it worked well for me +#idk if it'll work well for other people +class Feature(): + def __init__(self, alpha=0.1, decay=0.0005, curve=1, scaleType = 1): + self.min = None + self.max = None + self.alpha = alpha + self.decay = decay + self.last = 0.0 + self.curve = curve + self.scaleType = scaleType + + def update(self, x): + new = self.update_state(x) + self.last = self.last * self.alpha + new * (1 - self.alpha) + return self.last + + def update_state(self, x): + if self.min is None or self.max is None: + self.min = x - 0.00001 + self.max = x + 0.00001 + self.min = min(self.min, x) + self.max = max(self.max, x) + center = (self.min+self.max)/2 + + #The Min and Max decay slightly toward each other every frame allowing the range to change dynamically + self.min = (self.min * (1 - self.decay)) + (self.max * self.decay) + self.max = (self.max * (1 - self.decay)) + (self.min * self.decay) + + if self.scaleType == 1: + #Returns a value between -1 and 1 in relation to the maximum range + if x < center: + return -pow(clamp((x - center) / (self.min - center), 0, 1), self.curve) + elif x > center: + return pow(clamp((x - center) / (self.max - center), 0, 1), self.curve) + return 0 + if self.scaleType == 2: + #Returns a value between 0 and 1 in relation to the maximum range + return pow(clamp((x - self.min) / (self.max - self.min), 0, 1), self.curve) + return 0 + +class FeatureExtractor(): + def __init__(self): + self.eye_l = Feature(scaleType = 2, curve = 0.5) + self.eye_r = Feature(scaleType = 2, curve = 0.5) + self.eyebrow_updown_l = Feature( curve = 1.5) + self.eyebrow_updown_r = Feature( curve = 1.5) + self.mouth_corner_updown_l = Feature(curve = 2, decay=0.00005 ) + self.mouth_corner_updown_r = Feature(curve = 2, decay=0.00005 ) + self.mouth_open = Feature(scaleType = 2, curve = 1.5) + + def align_points(self, a, b, pts): #runs 6 times per frame + + alpha = (math.atan2(*(b - a)[::-1]) % (math.pi *2)) + + if alpha >= (math.pi/2): + alpha = - (alpha - math.pi) + if alpha <= -(math.pi/2): + alpha = - (alpha + math.pi) + aligned_pts = [] + for pt in pts: + aligned_pts.append(rotate(a, pt, alpha)) + return (aligned_pts) + + def update(self, pts): + features = {} + norm_distance_y = pts[27, 1]-pts[8, 1] + + f_pts = self.align_points(pts[42], pts[45], pts[[43, 44, 47, 46]]) + f = (f_pts[0][1]/2 + f_pts[1][1]/2) - (f_pts[2][1]/2 + f_pts[3][1]/2) + features["eye_l"] = self.eye_l.update(f) + f_pts = self.align_points(pts[36], pts[39], pts[[37, 38, 41, 40]]) + f = (f_pts[0][1]/2 + f_pts[1][1]/2) - (f_pts[2][1]/2 + f_pts[3][1]/2) + features["eye_r"] = self.eye_r.update(f) + + f = ((pts[22][1]/2 + pts[26][1]/2) - (pts[42][1]/2 + pts[45][1]/2))/norm_distance_y + features["eyebrow_updown_l"] = self.eyebrow_updown_l.update(f) + f = ((pts[17][1]/2 + pts[21][1]/2) - (pts[36][1]/2 + pts[39][1]/2))/norm_distance_y + features["eyebrow_updown_r"] = self.eyebrow_updown_r.update(f) + + f = ( pts[50][1] - pts[55][1]) * (norm_distance_y) + features["mouth_open"] = self.mouth_open.update(f)*0.66 + + #mouth corners are calculated together so they are assigned the same value + f = (((pts[51][1] + pts[55][1])/4) - ((pts[58][1]+pts[62][1])/2))/ norm_distance_y + features["mouth_corner_updown_r"] = self.mouth_corner_updown_r.update(f) * (1- 0.3 *features["mouth_open"]) * 0.66 - 0.1 + features["mouth_corner_updown_l"] = self.mouth_corner_updown_l.update(f) * (1- 0.3 * features["mouth_open"]) *0.66 - 0.1 + + #i removed some features VTS didn't seem to use + #the VTS library was handling them not being set anyway, so I just let that deal with the fact they're not here + return features diff --git a/landmarks.py b/landmarks.py new file mode 100644 index 0000000..8bb0f69 --- /dev/null +++ b/landmarks.py @@ -0,0 +1,266 @@ +import os +import numpy as np +import math +import cv2 +import onnxruntime +import time + +def matrix_to_quaternion(m): + t = 0.0 + q = [0.0, 0.0, 0, 0.0] + if m[2,2] < 0: + if m[0,0] > m[1,1]: + t = 1 + m[0,0] - m[1,1] - m[2,2] + q = [t, m[0,1]+m[1,0], m[2,0]+m[0,2], m[1,2]-m[2,1]] + else: + t = 1 - m[0,0] + m[1,1] - m[2,2] + q = [m[0,1]+m[1,0], t, m[1,2]+m[2,1], m[2,0]-m[0,2]] + else: + if m[0,0] < -m[1,1]: + t = 1 - m[0,0] - m[1,1] + m[2,2] + q = [m[2,0]+m[0,2], m[1,2]+m[2,1], t, m[0,1]-m[1,0]] + else: + t = 1 + m[0,0] + m[1,1] + m[2,2] + q = [m[1,2]-m[2,1], m[2,0]-m[0,2], m[0,1]-m[1,0], t] + return np.array(q) / (math.sqrt(t)* 0.5) + +#I haven't touched much here +#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 + + 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.face_3d = np.array([ + [ 0.4551769692672 , 0.300895790030204, -0.764429433974752], + [ 0.448998827123556, 0.166995837790733, -0.765143004071253], + [ 0.437431554952677, 0.022655479179981, -0.739267175112735], + [ 0.415033422928434, -0.088941454648772, -0.747947437846473], + [ 0.389123587370091, -0.232380029794684, -0.704788385327458], + [ 0.334630113904382, -0.361265387599081, -0.615587579236862], + [ 0.263725112132858, -0.460009725616771, -0.491479221041573], + [ 0.16241621322721 , -0.558037146073869, -0.339445180872282], + [ 0. , -0.621079019321682, -0.287294770748887], + [-0.16241621322721 , -0.558037146073869, -0.339445180872282], + [-0.263725112132858, -0.460009725616771, -0.491479221041573], + [-0.334630113904382, -0.361265387599081, -0.615587579236862], + [-0.389123587370091, -0.232380029794684, -0.704788385327458], + [-0.415033422928434, -0.088941454648772, -0.747947437846473], + [-0.437431554952677, 0.022655479179981, -0.739267175112735], + [-0.448998827123556, 0.166995837790733, -0.765143004071253], + [-0.4551769692672 , 0.300895790030204, -0.764429433974752], + [ 0.385529968662985, 0.402800553948697, -0.310031082540741], + [ 0.322196658344302, 0.464439136821772, -0.250558059367669], + [ 0.25409760441282 , 0.46420381416882 , -0.208177722146526], + [ 0.186875436782135, 0.44706071961879 , -0.145299823706503], + [ 0.120880983543622, 0.423566314072968, -0.110757158774771], + [-0.120880983543622, 0.423566314072968, -0.110757158774771], + [-0.186875436782135, 0.44706071961879 , -0.145299823706503], + [-0.25409760441282 , 0.46420381416882 , -0.208177722146526], + [-0.322196658344302, 0.464439136821772, -0.250558059367669], + [-0.385529968662985, 0.402800553948697, -0.310031082540741], + [ 0. , 0.293332603215811, -0.137582088779393], + [ 0. , 0.194828701837823, -0.069158109325951], + [ 0. , 0.103844017393155, -0.009151819844964], + [ 0. , 0. , 0. ], + [ 0.080626352317973, -0.041276068128093, -0.134161035564826], + [ 0.046439347377934, -0.057675223874769, -0.102990627164664], + [ 0. , -0.068753126205604, -0.090545348482397], + [-0.046439347377934, -0.057675223874769, -0.102990627164664], + [-0.080626352317973, -0.041276068128093, -0.134161035564826], + [ 0.315905195966084, 0.298337502555443, -0.285107407636464], + [ 0.275252345439353, 0.312721904921771, -0.244558251170671], + [ 0.176394511553111, 0.311907184376107, -0.219205360345231], + [ 0.131229723798772, 0.284447361805627, -0.234239149487417], + [ 0.184124948330084, 0.260179585304867, -0.226590776513707], + [ 0.279433549294448, 0.267363071770222, -0.248441437111633], + [-0.131229723798772, 0.284447361805627, -0.234239149487417], + [-0.176394511553111, 0.311907184376107, -0.219205360345231], + [-0.275252345439353, 0.312721904921771, -0.244558251170671], + [-0.315905195966084, 0.298337502555443, -0.285107407636464], + [-0.279433549294448, 0.267363071770222, -0.248441437111633], + [-0.184124948330084, 0.260179585304867, -0.226590776513707], + [ 0.121155252430729, -0.208988660580347, -0.160606287940521], + [ 0.041356305910044, -0.194484199722098, -0.096159882202821], + [ 0. , -0.205180167345702, -0.083299217789729], + [-0.041356305910044, -0.194484199722098, -0.096159882202821], + [-0.121155252430729, -0.208988660580347, -0.160606287940521], + [-0.132325402795928, -0.290857984604968, -0.187067868218105], + [-0.064137791831655, -0.325377847425684, -0.158924039726607], + [ 0. , -0.343742581679188, -0.113925986025684], + [ 0.064137791831655, -0.325377847425684, -0.158924039726607], + [ 0.132325402795928, -0.290857984604968, -0.187067868218105], + [ 0.181481567104525, -0.243239316141725, -0.231284988892766], + [ 0.083999507750469, -0.239717753728704, -0.155256465640701], + [ 0. , -0.256058040176369, -0.0950619498899 ], + [-0.083999507750469, -0.239717753728704, -0.155256465640701], + [-0.181481567104525, -0.243239316141725, -0.231284988892766], + [-0.074036069749345, -0.250689938345682, -0.177346470406188], + [ 0. , -0.264945854681568, -0.112349967428413], + [ 0.074036069749345, -0.250689938345682, -0.177346470406188], + # Pupils and eyeball centers + [ 0.257990002632141, 0.276080012321472, -0.219998998939991], + [-0.257990002632141, 0.276080012321472, -0.219998998939991], + [ 0.257990002632141, 0.276080012321472, -0.324570998549461], + [-0.257990002632141, 0.276080012321472, -0.324570998549461] + ], np.float32) + + self.contour_pts = [0,1,8,15,16,27,28,29,30,31,32,33,34,35] + 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 landmarks(self, tensor, crop_info): + crop_x1, crop_y1, scale_x, scale_y = crop_info + + t_main = tensor[0:66].reshape((66,784)) + t_m = t_main.argmax(1) + indices = np.expand_dims(t_m, 1) + + t_off_x = np.take_along_axis(tensor[66:132].reshape((66,784)), indices, 1).reshape((66,)) + p = np.clip(t_off_x, 0.0000001, 0.9999999) + t_off_x = 13.9375 * np.log(p / (1 - p)) + t_x = crop_y1 + scale_y * (223. * np.floor(t_m / 28) / 27.+ t_off_x) + + t_off_y = np.take_along_axis(tensor[132:198].reshape((66,784)), indices, 1).reshape((66,)) + p = np.clip(t_off_y, 0.0000001, 0.9999999) + t_off_y = 13.9375 * np.log(p / (1 - p)) + t_y = crop_x1 + scale_x * (223. * np.floor(np.mod(t_m, 28)) / 27. + t_off_y) + + t_conf = np.take_along_axis(t_main, indices, 1).reshape((66,)) + 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) + return (np.average(t_conf), lms) + + + + #I really need to break this guy up into smaller parts + 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.float32)), 0) + image_pts = np.array(lms)[self.contour_pts, 0:2] + + success = False + if face_info.rotation is not None: + success, face_info.rotation, face_info.translation = cv2.solvePnP(face_info.contour, image_pts, self.camera, np.zeros((4,1)), useExtrinsicGuess=True, rvec=np.transpose(face_info.rotation), tvec=np.transpose(face_info.translation), flags=cv2.SOLVEPNP_ITERATIVE) + else: + rvec = np.array([0, 0, 0], np.float32) + tvec = np.array([0, 0, 0], np.float32) + success, face_info.rotation, face_info.translation = cv2.solvePnP(face_info.contour, image_pts, self.camera, np.zeros((4,1)), useExtrinsicGuess=True, rvec=rvec, tvec=tvec, flags=cv2.SOLVEPNP_ITERATIVE) + + if not success: + face_info.rotation = np.array([0.0, 0.0, 0.0], np.float32) + face_info.translation = np.array([0.0, 0.0, 0.0], np.float32) + return False, np.zeros(4), np.zeros(3), 99999., pts_3d, lms + + rotation = face_info.rotation + translation = face_info.translation + pts_3d = np.zeros((70,3), np.float32) + + face_info.rotation = np.transpose(face_info.rotation) + face_info.translation = np.transpose(face_info.translation) + + rmat, _ = cv2.Rodrigues(rotation) + inverse_rotation = np.linalg.inv(rmat) + t_reference = face_info.face_3d.dot(rmat.transpose()) + t_reference = t_reference + face_info.translation + t_reference = t_reference.dot(self.camera.transpose()) + t_depth = t_reference[:, 2] + t_depth[t_depth == 0] = 0.000001 + t_depth_e = np.expand_dims(t_depth[:],1) + t_reference = t_reference[:] / t_depth_e + pts_3d[0:66] = np.stack([lms[0:66,0], lms[0:66,1], np.ones((66,))], 1) * t_depth_e[0:66] + pts_3d[0:66] = (pts_3d[0:66].dot(self.inverse_camera.transpose()) - face_info.translation).dot(inverse_rotation.transpose()) + + for i, pt in enumerate(face_info.face_3d[66:70]): + if i == 2: + # Right eyeball + # Eyeballs have an average diameter of 12.5mm and and the distance between eye corners is 30-35mm, so a conversion factor of 0.385 can be applied + eye_center = (pts_3d[36] + pts_3d[39]) / 2.0 + d_corner = np.linalg.norm(pts_3d[36] - pts_3d[39]) + depth = 0.385 * d_corner + pt_3d = np.array([eye_center[0], eye_center[1], eye_center[2] - depth]) + pts_3d[68] = pt_3d + continue + if i == 3: + # Left eyeball + eye_center = (pts_3d[42] + pts_3d[45]) / 2.0 + d_corner = np.linalg.norm(pts_3d[42] - pts_3d[45]) + depth = 0.385 * d_corner + pt_3d = np.array([eye_center[0], eye_center[1], eye_center[2] - depth]) + pts_3d[69] = pt_3d + continue + if i == 0: + d1 = np.linalg.norm(lms[66,0:2] - lms[36,0:2]) + d2 = np.linalg.norm(lms[66,0:2] - lms[39,0:2]) + d = d1 + d2 + pt = (pts_3d[36] * d1 + pts_3d[39] * d2) / d + if i == 1: + d1 = np.linalg.norm(lms[67,0:2] - lms[42,0:2]) + d2 = np.linalg.norm(lms[67,0:2] - lms[45,0:2]) + d = d1 + d2 + pt = (pts_3d[42] * d1 + pts_3d[45] * d2) / d + if i < 2: + reference = rmat.dot(pt) + reference = reference + face_info.translation + reference = self.camera.dot(reference) + depth = reference[2] + pt_3d = np.array([lms[66+i][0] * depth, lms[66+i][1] * depth, depth], np.float32) + pt_3d = self.inverse_camera.dot(pt_3d) + pt_3d = pt_3d - face_info.translation + pt_3d = inverse_rotation.dot(pt_3d) + pts_3d[66+i,:] = pt_3d[:] + + pts_3d[np.isnan(pts_3d).any(axis=1)] = np.array([0.,0.,0.], dtype=np.float32) + pnp_error = np.power(lms[0:17,0:2] - t_reference[0:17,0:2], 2).sum() + pnp_error += np.power(lms[30,0:2] - t_reference[30,0:2], 2).sum() + if np.isnan(pnp_error): + pnp_error = 9999999. + pnp_error = math.sqrt(pnp_error / (2.0 * image_pts.shape[0])) + if pnp_error > 300: + self.fail_count += 1 + if self.fail_count > 5: + # Something went wrong with adjusting the 3D model + print(f"Detected anomaly when 3D fitting face {face_info.id}. Resetting.") + face_info.face_3d = self.face_3d + face_info.rotation = None + face_info.translation = np.array([0.0, 0.0, 0.0], np.float32) + face_info.update_counts = np.zeros((66,2)) + face_info.update_contour() + else: + self.fail_count = 0 + + euler = cv2.RQDecomp3x3(rmat)[0] + return True, matrix_to_quaternion(rmat), euler, pnp_error, pts_3d, lms diff --git a/preview.py b/preview.py new file mode 100644 index 0000000..9c9da64 --- /dev/null +++ b/preview.py @@ -0,0 +1,21 @@ +import cv2 +import multiprocessing + +#this is it's own file just for ceanlieness + +def startProcess(frameQueue, targetFrameTime): + preview = Preview(frameQueue, targetFrameTime) + preview.start() + +class Preview(): + def __init__(self, frameQueue, targetFrameTime): + self.frameQueue = frameQueue + self.targetFrameTime = targetFrameTime + + def start(self): + while True: + #I had a timer on this, but tbh it's fine to just block it with the main thread + frame = self.frameQueue.get() + cv2.imshow("test",cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) + cv2.waitKey(1) + diff --git a/tracker.py b/tracker.py index f0588c2..561eb61 100644 --- a/tracker.py +++ b/tracker.py @@ -1,1040 +1,101 @@ -import os import numpy as np -import math import cv2 -import onnxruntime import time -import queue -import threading -import copy -from similaritytransform import SimilarityTransform -from retinaface import RetinaFaceDetector -from remedian import remedian +import facedetection +import eyes +import landmarks +import feature +import face -def resolve(name): - f = os.path.join(os.path.dirname(__file__), name) - return f +#this file is so much smaller than it used to be, lol +#I moved most of the functionality into separate files so they're easier to work with +def extract_face( frame, lms): + x1, y1 = lms.min(0) + x2, y2 = lms.max(0) + radius = np.array([(x2 - x1), (y2 - y1)])*0.6 + center = (np.array((x1, y1)) + np.array((x2, y2))) / 2.0 + w, h, _ = frame.shape + x1, y1 = clamp_to_im(center - radius, h, w) + x2, y2 = clamp_to_im(center + radius + 1, h, w) + offset = np.array((x1, y1)) + lms = (lms[:, 0:2] - offset).astype(int) + frame = frame[y1:y2, x1:x2] + return frame, lms, offset, center, radius -def clamp_to_im(pt, w, h): - x = pt[0] - y = pt[1] - if x < 0: - x = 0 - if y < 0: - y = 0 +def clamp_to_im(pt, w, h): #8 times per frame, but that only accounts for 0.005ms + x=max(pt[0],0) + y=max(pt[1],0) if x >= w: x = w-1 if y >= h: y = h-1 - return (int(x), int(y+1)) - -def rotate(origin, point, a): - a = -a - ox, oy = origin - px, py = point - - qx = ox + math.cos(a) * (px - ox) - math.sin(a) * (py - oy) - qy = oy + math.sin(a) * (px - ox) + math.cos(a) * (py - oy) - return qx, qy - -def angle(p1, p2): - p1 = np.array(p1) - p2 = np.array(p2) - a = np.arctan2(*(p2 - p1)[::-1]) - return (a % (2 * np.pi)) - -def compensate(p1, p2): - a = angle(p1, p2) - return rotate(p1, p2, a), a - -def rotate_image(image, a, center): - (h, w) = image.shape[:2] - a = np.rad2deg(a) - M = cv2.getRotationMatrix2D((float(center[0]), float(center[1])), a, 1.0) - rotated = cv2.warpAffine(image, M, (w, h)) - return rotated - -def intersects(r1, r2, amount=0.3): - area1 = r1[2] * r1[3] - area2 = r2[2] * r2[3] - inter = 0.0 - total = area1 + area2 - - r1_x1, r1_y1, w, h = r1 - r1_x2 = r1_x1 + w - r1_y2 = r1_y1 + h - r2_x1, r2_y1, w, h = r2 - r2_x2 = r2_x1 + w - r2_y2 = r2_y1 + h - - left = max(r1_x1, r2_x1) - right = min(r1_x2, r2_x2) - top = max(r1_y1, r2_y1) - bottom = min(r1_y2, r2_y2) - if left < right and top < bottom: - inter = (right - left) * (bottom - top) - total -= inter - - if inter / total >= amount: - return True - - return False - - #return not (r1_x1 > r2_x2 or r1_x2 < r2_x1 or r1_y1 > r2_y2 or r1_y2 < r2_y1) - -def group_rects(rects): - rect_groups = {} - for rect in rects: - rect_groups[str(rect)] = [-1, -1, []] - group_id = 0 - for i, rect in enumerate(rects): - name = str(rect) - group = group_id - group_id += 1 - if rect_groups[name][0] < 0: - rect_groups[name] = [group, -1, []] - else: - group = rect_groups[name][0] - for j, other_rect in enumerate(rects): - if i == j: - continue; - inter = intersects(rect, other_rect) - if intersects(rect, other_rect): - rect_groups[str(other_rect)] = [group, -1, []] - return rect_groups - -def logit(p, factor=16.0): - if p >= 1.0: - p = 0.9999999 - if p <= 0.0: - p = 0.0000001 - p = p/(1-p) - return float(np.log(p)) / float(factor) - -def logit_arr(p, factor=16.0): - p = np.clip(p, 0.0000001, 0.9999999) - return np.log(p / (1 - p)) / float(factor) - -def matrix_to_quaternion(m): - t = 0.0 - q = [0.0, 0.0, 0, 0.0] - if m[2,2] < 0: - if m[0,0] > m[1,1]: - t = 1 + m[0,0] - m[1,1] - m[2,2] - q = [t, m[0,1]+m[1,0], m[2,0]+m[0,2], m[1,2]-m[2,1]] - else: - t = 1 - m[0,0] + m[1,1] - m[2,2] - q = [m[0,1]+m[1,0], t, m[1,2]+m[2,1], m[2,0]-m[0,2]] - else: - if m[0,0] < -m[1,1]: - t = 1 - m[0,0] - m[1,1] + m[2,2] - q = [m[2,0]+m[0,2], m[1,2]+m[2,1], t, m[0,1]-m[1,0]] - else: - t = 1 + m[0,0] + m[1,1] + m[2,2] - q = [m[1,2]-m[2,1], m[2,0]-m[0,2], m[0,1]-m[1,0], t] - q = np.array(q, np.float32) * 0.5 / np.sqrt(t) - 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) - if conf > tracker.threshold: - try: - eye_state = tracker.get_eye_state(frame, lms) - except: - eye_state = [(1.0, 0.0, 0.0, 0.0), (1.0, 0.0, 0.0, 0.0)] - queue.put((session, conf, (lms, eye_state), crop_info, idx)) - else: - queue.put((session,)) - -class Feature(): - def __init__(self, threshold=0.15, alpha=0.2, hard_factor=0.15, decay=0.001, max_feature_updates=0): - self.median = remedian() - self.min = None - self.max = None - self.hard_min = None - self.hard_max = None - self.threshold = threshold - self.alpha = alpha - self.hard_factor = hard_factor - self.decay = decay - self.last = 0 - self.current_median = 0 - self.update_count = 0 - self.max_feature_updates = max_feature_updates - self.first_seen = -1 - self.updating = True - - def update(self, x, now=0): - if self.max_feature_updates > 0: - if self.first_seen == -1: - self.first_seen = now; - new = self.update_state(x, now=now) - filtered = self.last * self.alpha + new * (1 - self.alpha) - self.last = filtered - return filtered - - def update_state(self, x, now=0): - updating = self.updating and (self.max_feature_updates == 0 or now - self.first_seen < self.max_feature_updates) - if updating: - self.median + x - self.current_median = self.median.median() - else: - self.updating = False - median = self.current_median - - if self.min is None: - if x < median and (median - x) / median > self.threshold: - if updating: - self.min = x - self.hard_min = self.min + self.hard_factor * (median - self.min) - return -1 - return 0 - else: - if x < self.min: - if updating: - self.min = x - self.hard_min = self.min + self.hard_factor * (median - self.min) - return -1 - if self.max is None: - if x > median and (x - median) / median > self.threshold: - if updating: - self.max = x - self.hard_max = self.max - self.hard_factor * (self.max - median) - return 1 - return 0 - else: - if x > self.max: - if updating: - self.max = x - self.hard_max = self.max - self.hard_factor * (self.max - median) - return 1 - - if updating: - if self.min < self.hard_min: - self.min = self.hard_min * self.decay + self.min * (1 - self.decay) - if self.max > self.hard_max: - self.max = self.hard_max * self.decay + self.max * (1 - self.decay) - - if x < median: - return - (1 - (x - self.min) / (median - self.min)) - elif x > median: - return (x - median) / (self.max - median) - - return 0 - -class FeatureExtractor(): - def __init__(self, max_feature_updates=0): - self.eye_l = Feature(max_feature_updates=max_feature_updates) - self.eye_r = Feature(max_feature_updates=max_feature_updates) - self.eyebrow_updown_l = Feature(max_feature_updates=max_feature_updates) - self.eyebrow_updown_r = Feature(max_feature_updates=max_feature_updates) - self.eyebrow_quirk_l = Feature(threshold=0.05, max_feature_updates=max_feature_updates) - self.eyebrow_quirk_r = Feature(threshold=0.05, max_feature_updates=max_feature_updates) - self.eyebrow_steepness_l = Feature(threshold=0.05, max_feature_updates=max_feature_updates) - self.eyebrow_steepness_r = Feature(threshold=0.05, max_feature_updates=max_feature_updates) - self.mouth_corner_updown_l = Feature(max_feature_updates=max_feature_updates) - self.mouth_corner_updown_r = Feature(max_feature_updates=max_feature_updates) - self.mouth_corner_inout_l = Feature(threshold=0.02, max_feature_updates=max_feature_updates) - self.mouth_corner_inout_r = Feature(threshold=0.02, max_feature_updates=max_feature_updates) - self.mouth_open = Feature(max_feature_updates=max_feature_updates) - self.mouth_wide = Feature(threshold=0.02, max_feature_updates=max_feature_updates) - - def align_points(self, a, b, pts): - a = tuple(a) - b = tuple(b) - alpha = angle(a, b) - alpha = np.rad2deg(alpha) - if alpha >= 90: - alpha = - (alpha - 180) - if alpha <= -90: - alpha = - (alpha + 180) - alpha = np.deg2rad(alpha) - aligned_pts = [] - for pt in pts: - aligned_pts.append(np.array(rotate(a, pt, alpha))) - return alpha, np.array(aligned_pts) - - def update(self, pts, full=True): - features = {} - now = time.perf_counter() - - norm_distance_x = np.mean([pts[0, 0] - pts[16, 0], pts[1, 0] - pts[15, 0]]) - norm_distance_y = np.mean([pts[27, 1] - pts[28, 1], pts[28, 1] - pts[29, 1], pts[29, 1] - pts[30, 1]]) - - a1, f_pts = self.align_points(pts[42], pts[45], pts[[43, 44, 47, 46]]) - f = abs((np.mean([f_pts[0,1], f_pts[1,1]]) - np.mean([f_pts[2,1], f_pts[3,1]])) / norm_distance_y) - features["eye_l"] = self.eye_l.update(f, now) - - a2, f_pts = self.align_points(pts[36], pts[39], pts[[37, 38, 41, 40]]) - f = abs((np.mean([f_pts[0,1], f_pts[1,1]]) - np.mean([f_pts[2,1], f_pts[3,1]])) / norm_distance_y) - features["eye_r"] = self.eye_r.update(f, now) - - if full: - a3, _ = self.align_points(pts[0], pts[16], []) - 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. - - 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.mean([pts[17, 1], pts[21, 1]]) - pts[27, 1]) / norm_distance_y - features["eyebrow_updown_r"] = self.eyebrow_updown_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) - if full: - 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) - if full: - 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 = abs(np.mean(pts[[59,60,61], 1], axis=0) - np.mean(pts[[63,64,65], 1], axis=0)) / norm_distance_y - features["mouth_open"] = self.mouth_open.update(f, now) - - f = abs(pts[58, 0] - pts[62, 0]) / norm_distance_x - features["mouth_wide"] = self.mouth_wide.update(f, now) - - return features - -class FaceInfo(): - def __init__(self, id, tracker): - self.id = id - self.frame_count = -1 - self.tracker = tracker - self.contour_pts = [0,1,8,15,16,27,28,29,30,31,32,33,34,35] - self.face_3d = copy.copy(self.tracker.face_3d) - if self.tracker.model_type == -1: - self.contour_pts = [0,2,8,14,16,27,30,33] - self.reset() - self.alive = False - self.coord = None - self.base_scale_v = self.tracker.face_3d[27:30, 1] - self.tracker.face_3d[28:31, 1] - self.base_scale_h = np.abs(self.tracker.face_3d[[0, 36, 42], 0] - self.tracker.face_3d[[16, 39, 45], 0]) - - self.limit_3d_adjustment = True - self.update_count_delta = 75. - self.update_count_max = 7500. - - if self.tracker.max_feature_updates > 0: - self.features = FeatureExtractor(self.tracker.max_feature_updates) - - def reset(self): - self.alive = False - self.conf = None - self.lms = None - self.eye_state = None - self.rotation = None - self.translation = None - self.success = None - self.quaternion = None - self.euler = None - self.pnp_error = None - self.pts_3d = None - self.eye_blink = None - self.bbox = None - self.pnp_error = 0 - if self.tracker.max_feature_updates < 1: - self.features = FeatureExtractor(0) - self.current_features = {} - self.contour = np.zeros((21,3)) - self.update_counts = np.zeros((66,2)) - self.update_contour() - self.fail_count = 0 - - def update(self, result, coord, frame_count): - self.frame_count = frame_count - if result is None: - self.reset() - else: - self.conf, (self.lms, self.eye_state) = result - self.coord = coord - self.alive = True - - def update_contour(self): - self.contour = np.array(self.face_3d[self.contour_pts]) - - def normalize_pts3d(self, pts_3d): - # Calculate angle using nose - pts_3d[:, 0:2] -= pts_3d[30, 0:2] - alpha = angle(pts_3d[30, 0:2], pts_3d[27, 0:2]) - alpha -= np.deg2rad(90) - - R = np.matrix([[np.cos(alpha), -np.sin(alpha)], [np.sin(alpha), np.cos(alpha)]]) - pts_3d[:, 0:2] = (pts_3d - pts_3d[30])[:, 0:2].dot(R) + pts_3d[30, 0:2] - - # Vertical scale - pts_3d[:, 1] /= np.mean((pts_3d[27:30, 1] - pts_3d[28:31, 1]) / self.base_scale_v) - - # Horizontal scale - pts_3d[:, 0] /= np.mean(np.abs(pts_3d[[0, 36, 42], 0] - pts_3d[[16, 39, 45], 0]) / self.base_scale_h) - - return pts_3d - - def adjust_3d(self): - if self.conf < 0.4 or self.pnp_error > 300: - return - - 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 - update_type = -1 - d_o = np.ones((66,)) - d_c = np.ones((66,)) - for runs in range(max_runs): - r = 1.0 + np.random.random_sample((66,3)) * 0.02 - 0.01 - r[30, :] = 1.0 - if self.euler[0] > -165 and self.euler[0] < 145: - continue - elif self.euler[1] > -10 and self.euler[1] < 20: - r[:, 2] = 1.0 - update_type = 0 - else: - r[:, 0:2] = 1.0 - if self.euler[2] > 120 or self.euler[2] < 60: - continue - # Enable only one side of the points, depending on direction - elif self.euler[1] < -10: - update_type = 1 - r[[0, 1, 2, 3, 4, 5, 6, 7, 17, 18, 19, 20, 21, 31, 32, 36, 37, 38, 39, 40, 41, 48, 49, 56, 57, 58, 59, 65], 2] = 1.0 - eligible = [8, 9, 10, 11, 12, 13, 14, 15, 16, 22, 23, 24, 25, 26, 27, 28, 29, 33, 34, 35, 42, 43, 44, 45, 46, 47, 50, 51, 52, 53, 54, 55, 60, 61, 62, 63, 64] - else: - update_type = 1 - r[[9, 10, 11, 12, 13, 14, 15, 16, 22, 23, 24, 25, 26, 34, 35, 42, 43, 44, 45, 46, 47, 51, 52, 53, 54, 61, 62, 63], 2] = 1.0 - eligible = [0, 1, 2, 3, 4, 5, 6, 7, 8, 17, 18, 19, 20, 21, 27, 28, 29, 31, 32, 33, 36, 37, 38, 39, 40, 41, 48, 49, 50, 55, 56, 57, 58, 59, 60, 64, 65] - - if self.limit_3d_adjustment: - eligible = np.nonzero(self.update_counts[:, update_type] < self.update_counts[:, abs(update_type - 1)] + self.update_count_delta)[0] - if eligible.shape[0] <= 0: - break - - if runs == 0: - updated = copy.copy(self.face_3d[0:66]) - o_projected = np.ones((66,2)) - o_projected[eligible] = np.squeeze(np.array(cv2.projectPoints(self.face_3d[eligible], self.rotation, self.translation, self.tracker.camera, self.tracker.dist_coeffs)[0]), 1) - c = updated * r - c_projected = np.zeros((66,2)) - c_projected[eligible] = np.squeeze(np.array(cv2.projectPoints(c[eligible], self.rotation, self.translation, self.tracker.camera, self.tracker.dist_coeffs)[0]), 1) - changed = False - - d_o[eligible] = np.linalg.norm(o_projected[eligible] - self.lms[eligible, 0:2], axis=1) - d_c[eligible] = np.linalg.norm(c_projected[eligible] - self.lms[eligible, 0:2], axis=1) - indices = np.nonzero(d_c < d_o)[0] - if indices.shape[0] > 0: - if self.limit_3d_adjustment: - indices = np.intersect1d(indices, eligible) - if indices.shape[0] > 0: - self.update_counts[indices, update_type] += 1 - updated[indices] = c[indices] - o_projected[indices] = c_projected[indices] - changed = True - changed_any = changed_any or changed - - if not changed: - break - - if changed_any: - # Update weighted by point confidence - weights = np.zeros((66,3)) - weights[:, :] = self.lms[0:66, 2:3] - weights[weights > 0.7] = 1.0 - weights = 1.0 - weights - update_indices = np.arange(0, 66) - if self.limit_3d_adjustment: - update_indices = np.nonzero(self.update_counts[:, update_type] <= self.update_count_max)[0] - self.face_3d[update_indices] = self.face_3d[update_indices] * weights[update_indices] + updated[update_indices] * (1. - weights[update_indices]) - self.update_contour() - - self.pts_3d = self.normalize_pts3d(self.pts_3d) - if self.tracker.feature_level == 2: - self.current_features = self.features.update(self.pts_3d[:, 0:2]) - 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)) - 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)) - -def get_model_base_path(model_dir): - model_base_path = resolve(os.path.join("models")) - if model_dir is None: - if not os.path.exists(model_base_path): - model_base_path = resolve(os.path.join("..", "models")) - else: - model_base_path = model_dir - return model_base_path + return (int(x),int(y+1)) class Tracker(): - def __init__(self, width, height, model_type=3, detection_threshold=0.6, 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, try_hard=False): - options = onnxruntime.SessionOptions() - options.inter_op_num_threads = 1 - options.intra_op_num_threads = min(max_threads,4) - options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL - options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL - options.log_severity_level = 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 = "lm_modelT_opt.onnx" - if model_type >= 0: - model = self.models[self.model_type] - if model_type == -2: - model = "lm_modelV_opt.onnx" - if model_type == -3: - model = "lm_modelU_opt.onnx" - model_base_path = get_model_base_path(model_dir) + def __init__(self, width, height, messageQueue, model_type=3, detection_threshold=0.6, threshold=0.6, silent=False): - if threshold is None: - threshold = 0.6 - if model_type < 0: - threshold = 0.87 - - 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 - - # OTR 1.9 and later requires specifying providers - # explicitly as an argument in InferenceSession() - providersList = onnxruntime.capi._pybind_state.get_available_providers() - - # Single face instance with multiple threads - self.session = onnxruntime.InferenceSession(os.path.join(model_base_path, model), sess_options=options, providers=providersList) - - # Multiple faces with single threads - self.sessions = [] - self.max_workers = max(min(max_threads, max_faces), 1) - extra_threads = max_threads % self.max_workers - for i in range(self.max_workers): - options = onnxruntime.SessionOptions() - options.inter_op_num_threads = 1 - options.intra_op_num_threads = min(max(max_threads // self.max_workers, 4), 1) - if options.intra_op_num_threads < 1: - options.intra_op_num_threads = 1 - elif i < extra_threads: - options.intra_op_num_threads += 1 - options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL - options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL - self.sessions.append(onnxruntime.InferenceSession(os.path.join(model_base_path, model), sess_options=options, providers=providersList)) - self.input_name = self.session.get_inputs()[0].name - - 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) - - self.detection = onnxruntime.InferenceSession(os.path.join(model_base_path, "mnv3_detection_opt.onnx"), sess_options=options, providers=providersList) - self.faces = [] + self.FaceDetector = facedetection.FaceDetector(detection_threshold = detection_threshold) + self.EyeTracker = eyes.EyeTracker() + self.Landmarks = landmarks.Landmarks(width, height, model_type, threshold) + self.FeatureExtractor = feature.FeatureExtractor() # Image normalization constants - self.mean = np.float32(np.array([0.485, 0.456, 0.406])) - self.std = np.float32(np.array([0.229, 0.224, 0.225])) - 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], - [ 0.448998827123556, 0.166995837790733, -0.765143004071253], - [ 0.437431554952677, 0.022655479179981, -0.739267175112735], - [ 0.415033422928434, -0.088941454648772, -0.747947437846473], - [ 0.389123587370091, -0.232380029794684, -0.704788385327458], - [ 0.334630113904382, -0.361265387599081, -0.615587579236862], - [ 0.263725112132858, -0.460009725616771, -0.491479221041573], - [ 0.16241621322721 , -0.558037146073869, -0.339445180872282], - [ 0. , -0.621079019321682, -0.287294770748887], - [-0.16241621322721 , -0.558037146073869, -0.339445180872282], - [-0.263725112132858, -0.460009725616771, -0.491479221041573], - [-0.334630113904382, -0.361265387599081, -0.615587579236862], - [-0.389123587370091, -0.232380029794684, -0.704788385327458], - [-0.415033422928434, -0.088941454648772, -0.747947437846473], - [-0.437431554952677, 0.022655479179981, -0.739267175112735], - [-0.448998827123556, 0.166995837790733, -0.765143004071253], - [-0.4551769692672 , 0.300895790030204, -0.764429433974752], - [ 0.385529968662985, 0.402800553948697, -0.310031082540741], - [ 0.322196658344302, 0.464439136821772, -0.250558059367669], - [ 0.25409760441282 , 0.46420381416882 , -0.208177722146526], - [ 0.186875436782135, 0.44706071961879 , -0.145299823706503], - [ 0.120880983543622, 0.423566314072968, -0.110757158774771], - [-0.120880983543622, 0.423566314072968, -0.110757158774771], - [-0.186875436782135, 0.44706071961879 , -0.145299823706503], - [-0.25409760441282 , 0.46420381416882 , -0.208177722146526], - [-0.322196658344302, 0.464439136821772, -0.250558059367669], - [-0.385529968662985, 0.402800553948697, -0.310031082540741], - [ 0. , 0.293332603215811, -0.137582088779393], - [ 0. , 0.194828701837823, -0.069158109325951], - [ 0. , 0.103844017393155, -0.009151819844964], - [ 0. , 0. , 0. ], - [ 0.080626352317973, -0.041276068128093, -0.134161035564826], - [ 0.046439347377934, -0.057675223874769, -0.102990627164664], - [ 0. , -0.068753126205604, -0.090545348482397], - [-0.046439347377934, -0.057675223874769, -0.102990627164664], - [-0.080626352317973, -0.041276068128093, -0.134161035564826], - [ 0.315905195966084, 0.298337502555443, -0.285107407636464], - [ 0.275252345439353, 0.312721904921771, -0.244558251170671], - [ 0.176394511553111, 0.311907184376107, -0.219205360345231], - [ 0.131229723798772, 0.284447361805627, -0.234239149487417], - [ 0.184124948330084, 0.260179585304867, -0.226590776513707], - [ 0.279433549294448, 0.267363071770222, -0.248441437111633], - [-0.131229723798772, 0.284447361805627, -0.234239149487417], - [-0.176394511553111, 0.311907184376107, -0.219205360345231], - [-0.275252345439353, 0.312721904921771, -0.244558251170671], - [-0.315905195966084, 0.298337502555443, -0.285107407636464], - [-0.279433549294448, 0.267363071770222, -0.248441437111633], - [-0.184124948330084, 0.260179585304867, -0.226590776513707], - [ 0.121155252430729, -0.208988660580347, -0.160606287940521], - [ 0.041356305910044, -0.194484199722098, -0.096159882202821], - [ 0. , -0.205180167345702, -0.083299217789729], - [-0.041356305910044, -0.194484199722098, -0.096159882202821], - [-0.121155252430729, -0.208988660580347, -0.160606287940521], - [-0.132325402795928, -0.290857984604968, -0.187067868218105], - [-0.064137791831655, -0.325377847425684, -0.158924039726607], - [ 0. , -0.343742581679188, -0.113925986025684], - [ 0.064137791831655, -0.325377847425684, -0.158924039726607], - [ 0.132325402795928, -0.290857984604968, -0.187067868218105], - [ 0.181481567104525, -0.243239316141725, -0.231284988892766], - [ 0.083999507750469, -0.239717753728704, -0.155256465640701], - [ 0. , -0.256058040176369, -0.0950619498899 ], - [-0.083999507750469, -0.239717753728704, -0.155256465640701], - [-0.181481567104525, -0.243239316141725, -0.231284988892766], - [-0.074036069749345, -0.250689938345682, -0.177346470406188], - [ 0. , -0.264945854681568, -0.112349967428413], - [ 0.074036069749345, -0.250689938345682, -0.177346470406188], - # Pupils and eyeball centers - [ 0.257990002632141, 0.276080012321472, -0.219998998939991], - [-0.257990002632141, 0.276080012321472, -0.219998998939991], - [ 0.257990002632141, 0.276080012321472, -0.324570998549461], - [-0.257990002632141, 0.276080012321472, -0.324570998549461] - ], np.float32) + 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.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.dist_coeffs = np.zeros((4,1)) - - self.frame_count = 0 self.width = width self.height = height + self.threshold = threshold - self.detection_threshold = detection_threshold - self.max_faces = max_faces - self.max_threads = max_threads - self.discard = 0 - self.discard_after = discard_after - self.detected = 0 - self.wait_count = 0 - self.scan_every = scan_every - self.bbox_growth = bbox_growth - self.silent = silent - self.try_hard = try_hard - 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]) - if model_type < -1: - self.res = 112. - self.mean_res = np.tile(self.mean, [112, 112, 1]) - self.std_res = np.tile(self.std, [112, 112, 1]) - self.res_i = int(self.res) - self.out_res = 27. - if model_type < 0: - self.out_res = 6. - if model_type < -1: - self.out_res = 13. - self.out_res_i = int(self.out_res) + 1 - self.logit_factor = 16. - if model_type < 0: - self.logit_factor = 8. - if model_type < -1: - self.logit_factor = 16. + self.faces = [] + self.face_info = [face.FaceInfo(0, self)] + self.messageQueue = messageQueue - self.no_gaze = no_gaze - self.debug_gaze = False - self.feature_level = feature_level - if model_type == -1: - 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_224 + self.mean_224 + def preprocess(self, im): + im = cv2.resize(im, (224, 224), interpolation=cv2.INTER_CUBIC) * self.std + self.mean im = np.expand_dims(im, 0) - im = np.transpose(im, (0,3,1,2)) - outputs, maxpool = self.detection.run([], {'input': im}) - outputs = np.array(outputs) - maxpool = np.array(maxpool) - outputs[0, 0, outputs[0, 0] != maxpool[0, 0]] = 0 - detections = np.flip(np.argsort(outputs[0,0].flatten())) - results = [] - for det in detections[0:self.max_faces]: - y, x = det // 56, det % 56 - c = outputs[0, 0, y, x] - r = outputs[0, 1, y, x] * 112. - x *= 4 - y *= 4 - r *= 1.0 - if c < self.detection_threshold: - break - results.append((x - r, y - r, 2 * r, 2 * r * 1.0)) - results = np.array(results).astype(np.float32) - if results.shape[0] > 0: - results[:, [0,2]] *= frame.shape[1] / 224. - results[:, [1,3]] *= frame.shape[0] / 224. - return results + return np.transpose(im, (0,3,1,2)) - def landmarks(self, tensor, crop_info): - crop_x1, crop_y1, scale_x, scale_y, _ = crop_info - avg_conf = 0 - res = self.res - 1 - c0, c1, c2 = 66, 132, 198 - if self.model_type == -1: - 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((c0,)) - 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[c1:c2].reshape((c0,self.out_res_i * self.out_res_i)), indices, 1).reshape((c0,)) - t_off_x = res * logit_arr(t_off_x, self.logit_factor) - 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 == -1: - 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]] - #lms[[1,3,4,6,7,9,10,12,13,15,18,20,23,25,38,40,44,46]] += lms[[2,2,5,5,8,8,11,11,14,16,19,21,24,26,39,39,45,45]] - #lms[[3,4,6,7,9,10,12,13]] += lms[[5,5,8,8,11,11,14,14]] - #lms[[1,15,18,20,23,25,38,40,44,46]] /= 2.0 - #lms[[3,4,6,7,9,10,12,13]] /= 3.0 - 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 cropFace(self,face, im): + duration_pp = 0.0 + x,y,w,h = face - 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.float32)), 0) + wint = int(w * 0.1) + hint = int(h * 0.125) + crop_x1 = x - wint + crop_y1 = y - hint + crop_x2 = x + w + wint + crop_y2 = y + h + hint - image_pts = np.array(lms)[face_info.contour_pts, 0:2] + crop_x1, crop_y1 = clamp_to_im((crop_x1, crop_y1), self.width, self.height) + crop_x2, crop_y2 = clamp_to_im((crop_x2, crop_y2), self.width, self.height) + cropx = (crop_x2 - crop_x1) + cropy = (crop_y2 - crop_y1) - success = False - if face_info.rotation is not None: - success, face_info.rotation, face_info.translation = cv2.solvePnP(face_info.contour, image_pts, self.camera, self.dist_coeffs, useExtrinsicGuess=True, rvec=np.transpose(face_info.rotation), tvec=np.transpose(face_info.translation), flags=cv2.SOLVEPNP_ITERATIVE) - else: - rvec = np.array([0, 0, 0], np.float32) - tvec = np.array([0, 0, 0], np.float32) - success, face_info.rotation, face_info.translation = cv2.solvePnP(face_info.contour, image_pts, self.camera, self.dist_coeffs, useExtrinsicGuess=True, rvec=rvec, tvec=tvec, flags=cv2.SOLVEPNP_ITERATIVE) + if cropx < 64 or cropy < 64: + return (None, [], duration_pp ) - rotation = face_info.rotation - translation = face_info.translation + start_pp = time.perf_counter() + crop = self.preprocess(im[crop_y1:crop_y2, crop_x1:crop_x2]) + duration_pp += 1000 * (time.perf_counter() - start_pp) - pts_3d = np.zeros((70,3), np.float32) - if not success: - face_info.rotation = np.array([0.0, 0.0, 0.0], np.float32) - face_info.translation = np.array([0.0, 0.0, 0.0], np.float32) - return False, np.zeros(4), np.zeros(3), 99999., pts_3d, lms - else: - face_info.rotation = np.transpose(face_info.rotation) - face_info.translation = np.transpose(face_info.translation) + scale_x = cropx / 224. + scale_y = cropy / 224. + crop_info = (crop_x1, crop_y1, scale_x, scale_y) + return (crop, crop_info, duration_pp ) - rmat, _ = cv2.Rodrigues(rotation) - inverse_rotation = np.linalg.inv(rmat) - t_reference = face_info.face_3d.dot(rmat.transpose()) - t_reference = t_reference + face_info.translation - t_reference = t_reference.dot(self.camera.transpose()) - t_depth = t_reference[:, 2] - t_depth[t_depth == 0] = 0.000001 - t_depth_e = np.expand_dims(t_depth[:],1) - t_reference = t_reference[:] / t_depth_e - pts_3d[0:66] = np.stack([lms[0:66,0], lms[0:66,1], np.ones((66,))], 1) * t_depth_e[0:66] - pts_3d[0:66] = (pts_3d[0:66].dot(self.inverse_camera.transpose()) - face_info.translation).dot(inverse_rotation.transpose()) - pnp_error = np.power(lms[0:17,0:2] - t_reference[0:17,0:2], 2).sum() - pnp_error += np.power(lms[30,0:2] - t_reference[30,0:2], 2).sum() - if np.isnan(pnp_error): - pnp_error = 9999999. - for i, pt in enumerate(face_info.face_3d[66:70]): - if i == 2: - # Right eyeball - # Eyeballs have an average diameter of 12.5mm and and the distance between eye corners is 30-35mm, so a conversion factor of 0.385 can be applied - eye_center = (pts_3d[36] + pts_3d[39]) / 2.0 - d_corner = np.linalg.norm(pts_3d[36] - pts_3d[39]) - depth = 0.385 * d_corner - pt_3d = np.array([eye_center[0], eye_center[1], eye_center[2] - depth]) - pts_3d[68] = pt_3d - continue - if i == 3: - # Left eyeball - eye_center = (pts_3d[42] + pts_3d[45]) / 2.0 - d_corner = np.linalg.norm(pts_3d[42] - pts_3d[45]) - depth = 0.385 * d_corner - pt_3d = np.array([eye_center[0], eye_center[1], eye_center[2] - depth]) - pts_3d[69] = pt_3d - continue - if i == 0: - d1 = np.linalg.norm(lms[66,0:2] - lms[36,0:2]) - d2 = np.linalg.norm(lms[66,0:2] - lms[39,0:2]) - d = d1 + d2 - pt = (pts_3d[36] * d1 + pts_3d[39] * d2) / d - if i == 1: - d1 = np.linalg.norm(lms[67,0:2] - lms[42,0:2]) - d2 = np.linalg.norm(lms[67,0:2] - lms[45,0:2]) - d = d1 + d2 - pt = (pts_3d[42] * d1 + pts_3d[45] * d2) / d - if i < 2: - reference = rmat.dot(pt) - reference = reference + face_info.translation - reference = self.camera.dot(reference) - depth = reference[2] - pt_3d = np.array([lms[66+i][0] * depth, lms[66+i][1] * depth, depth], np.float32) - pt_3d = self.inverse_camera.dot(pt_3d) - pt_3d = pt_3d - face_info.translation - pt_3d = inverse_rotation.dot(pt_3d) - pts_3d[66+i,:] = pt_3d[:] - pts_3d[np.isnan(pts_3d).any(axis=1)] = np.array([0.,0.,0.], dtype=np.float32) + def early_exit(self, reason, start): + self.messageQueue.put(reason) + self.faces = [] + duration = (time.perf_counter() - start) * 1000 + self.messageQueue.put(f"Took {duration:.2f}ms") + return None, None, None - pnp_error = np.sqrt(pnp_error / (2.0 * image_pts.shape[0])) - if pnp_error > 300: - face_info.fail_count += 1 - if face_info.fail_count > 5: - # Something went wrong with adjusting the 3D model - if not self.silent: - print(f"Detected anomaly when 3D fitting face {face_info.id}. Resetting.") - face_info.face_3d = copy.copy(self.face_3d) - face_info.rotation = None - face_info.translation = np.array([0.0, 0.0, 0.0], np.float32) - face_info.update_counts = np.zeros((66,2)) - face_info.update_contour() - else: - face_info.fail_count = 0 - - euler = cv2.RQDecomp3x3(rmat)[0] - return True, matrix_to_quaternion(rmat), euler, pnp_error, pts_3d, lms - - 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_res + self.mean_res - im = np.expand_dims(im, 0) - im = np.transpose(im, (0,3,1,2)) - return im - - def equalize(self, im): - im_yuv = cv2.cvtColor(im, cv2.COLOR_BGR2YUV) - im_yuv[:,:,0] = cv2.equalizeHist(im_yuv[:,:,0]) - return cv2.cvtColor(im_yuv, cv2.COLOR_YUV2BGR) - - def corners_to_eye(self, corners, w, h, flip): - ((cx1, cy1), (cx2, cy2)) = corners - c1 = np.array([cx1, cy1]) - c2 = np.array([cx2, cy2]) - c2, a = compensate(c1, c2) - center = (c1 + c2) / 2.0 - radius = np.linalg.norm(c1 - c2) / 2.0 - radius = np.array([radius * 1.4, radius * 1.2]) - upper_left = clamp_to_im(center - radius, w, h) - lower_right = clamp_to_im(center + radius, w, h) - return upper_left, lower_right, center, radius, c1, a - - def prepare_eye(self, frame, full_frame, lms, flip): - outer_pt = tuple(lms[0]) - inner_pt = tuple(lms[1]) - h, w, _ = frame.shape - (x1, y1), (x2, y2), center, radius, reference, a = self.corners_to_eye((outer_pt, inner_pt), w, h, flip) - im = rotate_image(frame[:, :, ::], a, reference) - im = im[int(y1):int(y2), int(x1):int(x2),:] - if np.prod(im.shape) < 1: - return None, None, None, None, None, None - if flip: - im = cv2.flip(im, 1) - scale = np.array([(x2 - x1), (y2 - y1)]) / 32. - im = cv2.resize(im, (32, 32), interpolation=cv2.INTER_LINEAR) - #im = self.equalize(im) - if self.debug_gaze: - if not flip: - full_frame[0:32, 0:32] = im - else: - full_frame[0:32, 32:64] = im - im = im.astype(np.float32)[:,:,::-1] * self.std_32 + self.mean_32 - im = np.expand_dims(im, 0) - im = np.transpose(im, (0,3,2,1)) - return im, x1, y1, scale, reference, a - - def extract_face(self, frame, lms): - lms = np.array(lms)[:,0:2][:,::-1] - x1, y1 = tuple(lms.min(0)) - x2, y2 = tuple(lms.max(0)) - radius_x = 1.2 * (x2 - x1) / 2.0 - radius_y = 1.2 * (y2 - y1) / 2.0 - radius = np.array((radius_x, radius_y)) - center = (np.array((x1, y1)) + np.array((x2, y2))) / 2.0 - w, h, _ = frame.shape - x1, y1 = clamp_to_im(center - radius, h, w) - x2, y2 = clamp_to_im(center + radius + 1, h, w) - offset = np.array((x1, y1)) - lms = (lms[:, 0:2] - offset).astype(int) - frame = frame[y1:y2, x1:x2] - return frame, lms, offset - - def get_eye_state(self, frame, lms): - if self.no_gaze: - return [(1.0, 0.0, 0.0, 0.0), (1.0, 0.0, 0.0, 0.0)] - lms = np.array(lms) - e_x = [0,0] - e_y = [0,0] - scale = [0,0] - reference = [None, None] - angles = [0, 0] - face_frame, lms, offset = self.extract_face(frame, lms) - (right_eye, e_x[0], e_y[0], scale[0], reference[0], angles[0]) = self.prepare_eye(face_frame, frame, np.array([lms[36,0:2], lms[39,0:2]]), False) - (left_eye, e_x[1], e_y[1], scale[1], reference[1], angles[1]) = self.prepare_eye(face_frame, frame, np.array([lms[42,0:2], lms[45,0:2]]), True) - if right_eye is None or left_eye is None: - return [(1.0, 0.0, 0.0, 0.0), (1.0, 0.0, 0.0, 0.0)] - both_eyes = np.concatenate((right_eye, left_eye)) - results = None - results = self.gaze_model.run([], {self.input_name: both_eyes}) - open = [0, 0] - open[0] = 1#results[1][0].argmax() - open[1] = 1#results[1][1].argmax() - results = np.array(results[0]) - - eye_state = [] - for i in range(2): - m = int(results[i][0].argmax()) - x = m // 8 - y = m % 8 - conf = float(results[i][0][x,y]) - - off_x = 32.0 * logit(results[i][1][x, y], 8.0) - off_y = 32.0 * logit(results[i][2][x, y], 8.0) - if i == 1: - eye_x = 32.0 * float(x) / 8.0 + off_x - else: - eye_x = 32.0 * float(x) / 8.0 + off_x - eye_y = 32.0 * float(y) / 8.0 + off_y - - if self.debug_gaze: - if i == 0: - frame[int(eye_y), int(eye_x)] = (0, 0, 255) - frame[int(eye_y+1), int(eye_x)] = (0, 0, 255) - frame[int(eye_y+1), int(eye_x+1)] = (0, 0, 255) - frame[int(eye_y), int(eye_x+1)] = (0, 0, 255) - else: - frame[int(eye_y), 32+int(eye_x)] = (0, 0, 255) - frame[int(eye_y+1), 32+int(eye_x)] = (0, 0, 255) - frame[int(eye_y+1), 32+int(eye_x+1)] = (0, 0, 255) - frame[int(eye_y), 32+int(eye_x+1)] = (0, 0, 255) - - if i == 0: - eye_x = e_x[i] + scale[i][0] * eye_x - else: - eye_x = e_x[i] + scale[i][0] * (32. - eye_x) - eye_y = e_y[i] + scale[i][1] * eye_y - eye_x, eye_y = rotate(reference[i], (eye_x, eye_y), -angles[i]) - - eye_x = eye_x + offset[0] - eye_y = eye_y + offset[1] - eye_state.append([open[i], eye_y, eye_x, conf]) - - eye_state = np.array(eye_state) - eye_state[np.isnan(eye_state).any(axis=1)] = np.array([1.,0.,0.,0.], dtype=np.float32) - 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: - adjusted_results.append((conf - conf_adjust, (lms, eye_state))) - result_coords.append(np.array(lms)[:, 0:2].mean(0)) - results = adjusted_results - candidates = [[]] * self.max_faces - max_dist = 2 * np.linalg.norm(np.array([self.width, self.height])) - for i, face_info in enumerate(self.face_info): - for j, coord in enumerate(result_coords): - if face_info.coord is None: - candidates[i].append((max_dist, i, j)) - else: - candidates[i].append((np.linalg.norm(face_info.coord - coord), i, j)) - for i, candidate in enumerate(candidates): - candidates[i] = sorted(candidate) - found = 0 - target = len(results) - used_results = {} - used_faces = {} - while found < target: - min_list = min(candidates) - candidate = min_list.pop(0) - face_idx = candidate[1] - result_idx = candidate[2] - if not result_idx in used_results and not face_idx in used_faces: - self.face_info[face_idx].update(results[result_idx], result_coords[result_idx], self.frame_count) - min_list.clear() - used_results[result_idx] = True - used_faces[face_idx] = True - found += 1 - if len(min_list) == 0: - min_list.append((2 * max_dist, face_idx, result_idx)) - for face_info in self.face_info: - if face_info.frame_count != self.frame_count: - face_info.update(None, None, self.frame_count) - - def predict(self, frame, additional_faces=[]): - self.frame_count += 1 + def predict(self, frame, face = None): start = time.perf_counter() - im = frame - duration_fd = 0.0 duration_pp = 0.0 duration_model = 0.0 @@ -1042,170 +103,53 @@ class Tracker(): new_faces = [] new_faces.extend(self.faces) - bonus_cutoff = len(self.faces) - new_faces.extend(additional_faces) - self.wait_count += 1 - if self.detected == 0: - start_fd = time.perf_counter() - if self.use_retinaface > 0 or self.try_hard: - retinaface_detections = self.retinaface.detect_retina(frame) - new_faces.extend(retinaface_detections) - if self.use_retinaface == 0 or self.try_hard: - new_faces.extend(self.detect_faces(frame)) - if self.try_hard: - new_faces.extend([(0, 0, self.width, self.height)]) - duration_fd = 1000 * (time.perf_counter() - start_fd) - self.wait_count = 0 - elif self.detected < self.max_faces: - if self.use_retinaface > 0: - new_faces.extend(self.retinaface_scan.get_results()) - if self.wait_count >= self.scan_every: - if self.use_retinaface > 0: - self.retinaface_scan.background_detect(frame) - else: - start_fd = time.perf_counter() - new_faces.extend(self.detect_faces(frame)) - duration_fd = 1000 * (time.perf_counter() - start_fd) - self.wait_count = 0 - else: - self.wait_count = 0 + if face is not None: + new_faces.append(face) if len(new_faces) < 1: - duration = (time.perf_counter() - start) * 1000 - if not self.silent: - print(f"Took {duration:.2f}ms") - return [] + start_fd = time.perf_counter() + new_faces.extend(self.FaceDetector.detect_faces(frame)) + duration_fd = 1000 * (time.perf_counter() - start_fd) + if len(new_faces) < 1: + return self.early_exit("No faces found", start) - crops = [] - crop_info = [] - num_crops = 0 - for j, (x,y,w,h) in enumerate(new_faces): - crop_x1 = x - int(w * 0.1) - crop_y1 = y - int(h * 0.125) - crop_x2 = x + w + int(w * 0.1) - crop_y2 = y + h + int(h * 0.125) + crop, crop_info, duration_pp = self.cropFace(new_faces[0], frame) - crop_x1, crop_y1 = clamp_to_im((crop_x1, crop_y1), self.width, self.height) - crop_x2, crop_y2 = clamp_to_im((crop_x2, crop_y2), self.width, self.height) - - scale_x = float(crop_x2 - crop_x1) / self.res - scale_y = float(crop_y2 - crop_y1) / self.res - - if crop_x2 - crop_x1 < 4 or crop_y2 - crop_y1 < 4: - continue - - start_pp = time.perf_counter() - crop = self.preprocess(im, (crop_x1, crop_y1, crop_x2, crop_y2)) - duration_pp += 1000 * (time.perf_counter() - start_pp) - crops.append(crop) - crop_info.append((crop_x1, crop_y1, scale_x, scale_y, 0.0 if j >= bonus_cutoff else 0.1)) - num_crops += 1 + #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() - outputs = {} - if num_crops == 1: - output = self.session.run([], {self.input_name: crops[0]})[0] - conf, lms = self.landmarks(output[0], crop_info[0]) - if conf > self.threshold: - try: - eye_state = self.get_eye_state(frame, lms) - except: - eye_state = [(1.0, 0.0, 0.0, 0.0), (1.0, 0.0, 0.0, 0.0)] - outputs[crop_info[0]] = (conf, (lms, eye_state), 0) - else: - started = 0 - results = queue.Queue() - for i in range(min(num_crops, self.max_workers)): - thread = threading.Thread(target=worker_thread, args=(self.sessions[started], frame, crops[started], crop_info[started], results, self.input_name, started, self)) - started += 1 - thread.start() - returned = 0 - while returned < num_crops: - result = results.get(True) - if len(result) != 1: - session, conf, lms, sample_crop_info, idx = result - outputs[sample_crop_info] = (conf, lms, idx) - else: - session = result[0] - returned += 1 - if started < num_crops: - thread = threading.Thread(target=worker_thread, args=(session, frame, crops[started], crop_info[started], results, self.input_name, started, self)) - started += 1 - thread.start() + #self.Landmarks - actual_faces = [] - good_crops = [] - for crop in crop_info: - if crop not in outputs: - continue - conf, lms, i = outputs[crop] - x1, y1, _ = lms[0].min(0) - x2, y2, _ = lms[0].max(0) - bb = (x1, y1, x2 - x1, y2 - y1) - outputs[crop] = (conf, lms, i, bb) - actual_faces.append(bb) - good_crops.append(crop) - groups = group_rects(actual_faces) + conf, lms = self.Landmarks.run(crop , crop_info ) + #Early exit if below confidence threshold + if conf < self.threshold: + return self.early_exit("Confidence below threshold", start) - best_results = {} - for crop in good_crops: - conf, lms, i, bb = outputs[crop] - if conf < self.threshold: - continue; - group_id = groups[str(bb)][0] - if not group_id in best_results: - best_results[group_id] = [-1, [], 0] - if conf > self.threshold and best_results[group_id][0] < conf + crop[4]: - best_results[group_id][0] = conf + crop[4] - best_results[group_id][1] = lms - best_results[group_id][2] = crop[4] + #lms is short for landmarks + face_frame, face_lms, face_offset,face_center, face_radius = extract_face(frame, np.array(lms)[:,0:2][:,::-1]) + eye_state = self.EyeTracker.get_eye_state(face_frame, face_lms, face_offset) + + self.face_info[0].update((conf, (lms, eye_state)), np.array(lms)[:, 0:2].mean(0)) - sorted_results = sorted(best_results.values(), key=lambda x: x[0], reverse=True)[:self.max_faces] - self.assign_face_info(sorted_results) duration_model = 1000 * (time.perf_counter() - start_model) - - results = [] - detected = [] start_pnp = time.perf_counter() - for face_info in self.face_info: - if face_info.alive and face_info.conf > self.threshold: - face_info.success, face_info.quaternion, face_info.euler, face_info.pnp_error, face_info.pts_3d, face_info.lms = self.estimate_depth(face_info) + face_info = self.face_info[0] + + 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) + + if face_info.success: face_info.adjust_3d() lms = face_info.lms[:, 0:2] - x1, y1 = tuple(lms[0:66].min(0)) - x2, y2 = tuple(lms[0:66].max(0)) - bbox = (y1, x1, y2 - y1, x2 - x1) - face_info.bbox = bbox - detected.append(bbox) - results.append(face_info) - duration_pnp += 1000 * (time.perf_counter() - start_pnp) + x1, y1 = lms[0:66].min(0) + x2, y2 = lms[0:66].max(0) + self.faces = [[y1, x1, y2 - y1, x2 - x1]] + duration_pnp += 1000 * (time.perf_counter() - start_pnp) + duration = (time.perf_counter() - start) * 1000 + self.messageQueue.put(f"Took {duration:.2f}ms (detect: {duration_fd:.2f}ms, crop: {duration_pp:.2f}ms, track: {duration_model:.2f}ms, 3D points: {duration_pnp:.2f}ms)") + return face_info, face_center, face_radius - if len(detected) > 0: - self.detected = len(detected) - self.faces = detected - self.discard = 0 - else: - self.detected = 0 - self.discard += 1 - if self.discard > self.discard_after: - self.faces = [] - else: - if self.bbox_growth > 0: - faces = [] - for (x,y,w,h) in self.faces: - x -= w * self.bbox_growth - y -= h * self.bbox_growth - w += 2 * w * self.bbox_growth - h += 2 * h * self.bbox_growth - faces.append((x,y,w,h)) - self.faces = faces - self.faces = [x for x in self.faces if not np.isnan(np.array(x)).any()] - self.detected = len(self.faces) - - duration = (time.perf_counter() - start) * 1000 - if not self.silent: - print(f"Took {duration:.2f}ms (detect: {duration_fd:.2f}ms, crop: {duration_pp:.2f}ms, track: {duration_model:.2f}ms, 3D points: {duration_pnp:.2f}ms)") - - results = sorted(results, key=lambda x: x.id) - - return results + #Combined multiple failures into one catch all exit + return self.early_exit("Face info not valid", start) diff --git a/vts.py b/vts.py new file mode 100644 index 0000000..b041ac2 --- /dev/null +++ b/vts.py @@ -0,0 +1,72 @@ +import time +import socket +import struct + +class VTS(): + def __init__(self, target_ip, target_port, silent, height, width, packetQueue, messageQueue): + + self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self.target_ip = target_ip + self.target_port = target_port + self.silent = silent + self.features = ["eye_l", "eye_r", "eyebrow_steepness_l", "eyebrow_updown_l", "eyebrow_quirk_l", "eyebrow_steepness_r", "eyebrow_updown_r", "eyebrow_quirk_r", "mouth_corner_updown_l", "mouth_corner_inout_l", "mouth_corner_updown_r", "mouth_corner_inout_r", "mouth_open", "mouth_wide"] + self.width = width + self.height = height + self.messageQueue = messageQueue + self.packetQueue = packetQueue + + def packetSender(self): + while True: + packet = self.packetQueue.get() + self.sendPacket(packet) + + def preparePacket(self, face): + packet = bytearray() + if face.eye_blink is None: + face.eye_blink = [1, 1] + if not self.silent: + right_state = "O" if face.eye_blink[0] > 0.2 else "-" + left_state = "O" if face.eye_blink[1] > 0.2 else "-" + 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("f", self.width))) + packet.extend(bytearray(struct.pack("f", self.height))) + packet.extend(bytearray(struct.pack("f", face.eye_blink[0]))) + packet.extend(bytearray(struct.pack("f", face.eye_blink[1]))) + packet.extend(bytearray(struct.pack("B", 1))) + packet.extend(bytearray(struct.pack("f", face.pnp_error))) + packet.extend(bytearray(struct.pack("f", face.quaternion[0]))) + packet.extend(bytearray(struct.pack("f", face.quaternion[1]))) + packet.extend(bytearray(struct.pack("f", face.quaternion[2]))) + packet.extend(bytearray(struct.pack("f", face.quaternion[3]))) + packet.extend(bytearray(struct.pack("f", face.euler[0]))) + packet.extend(bytearray(struct.pack("f", face.euler[1]))) + packet.extend(bytearray(struct.pack("f", face.euler[2]))) + packet.extend(bytearray(struct.pack("f", face.translation[0]))) + packet.extend(bytearray(struct.pack("f", face.translation[1]))) + packet.extend(bytearray(struct.pack("f", face.translation[2]))) + for (_,_,c) in face.lms: + packet.extend(bytearray(struct.pack("f", c))) + for pt_num, (x,y,_) in enumerate(face.lms): + packet.extend(bytearray(struct.pack("f", y))) + packet.extend(bytearray(struct.pack("f", x))) + for (x,y,z) in face.pts_3d: + packet.extend(bytearray(struct.pack("f", x))) + packet.extend(bytearray(struct.pack("f", -y))) + packet.extend(bytearray(struct.pack("f", -z))) + if face.current_features is None: + face.current_features = {} + for feature in self.features: + if not feature in face.current_features: + face.current_features[feature] = 0 + packet.extend(bytearray(struct.pack("f", face.current_features[feature]))) + return packet + + def sendPacket(self, packet): + try: + self.sock.sendto(packet, (self.target_ip, self.target_port)) + except: + self.messageQueue.put("Failed to send packet") + return diff --git a/webcam.py b/webcam.py new file mode 100644 index 0000000..d56138b --- /dev/null +++ b/webcam.py @@ -0,0 +1,104 @@ +import cv2 +import multiprocessing +import time +import numpy as np +import math + + +#this is just here for conveience +def startProcess(frameQueue, faceQueue, fps, targetBrightness, width, height, mirrorInput, messageQueue): + webcam = Webcam(frameQueue, faceQueue, fps, targetBrightness, width, height, mirrorInput, messageQueue) + webcam.start() + +class Webcam(): + def __init__(self, frameQueue, faceQueue, fps, targetBrightness, width, height, mirrorInput, messageQueue): + self.cap = cv2.VideoCapture(0) + #tbh, these settings will make things way slower and frametimes will suffer + #but they're the only way frames bigger than 480p will work + if height > 480: + self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*"MJPG")) + self.cap.set(38, 4) + else: + self.cap.set(38, 2) + self.cap.set(3, width) + self.cap.set(4, height) + self.cap.set(cv2.CAP_PROP_FPS, fps) + self.targetFrameTime = 1. / (fps-0.01) + self.gamma = 0.7 + self.mirror = mirrorInput + self.frameQueue = frameQueue #outgoing frames from the webcam + self.faceQueue = faceQueue #incoming face data for gamma calculation + self.cameraLatency = 0 + self.brightnessFrame = None #a copy of the brightness channel for gamma calculations + self.targetBrightness = targetBrightness #the target average brightness of the face, used in gamma calculations + self.messageQueue = messageQueue #outgoing queue, unused, but I wanted it here in case I wanted it + self.width = width #unused, but it seemed useful to have around + self.height = height #unused, but it seemed useful to have around + self.frame = None + self.ret = 0 + + def start(self): + while self.cap.isOpened(): + frameStart= time.perf_counter() + self.getFrame() + self.cameraLatency = time.perf_counter() - frameStart + if self.ret: + self.applyGamma() + if self.mirror: + self.frame = cv2.flip(self.frame, 1) + #this line is where this process sits 90% of the time + self.frameQueue.put([self.frame, self.cameraLatency]) + if self.faceQueue.qsize() > 0: + self.updateGamma() + sleepTime = self.targetFrameTime - (time.perf_counter() - frameStart) + sleepTime = max(sleepTime, 0) + time.sleep(sleepTime) + + def getFrame(self): + #My webcam's gain tends to be noisy + #and my gamma correction seems less noisy + #so I tell the camera not to apply gain + self.cap.set(cv2.CAP_PROP_GAIN, 0) + #keeping track of how long it takes to get frames from the webcam + #because that one line is the cause of 90% of late frames + cameraStart = time.perf_counter() + self.ret, self.frame = self.cap.read() + self.cameraLatency = time.perf_counter() - cameraStart + return + + #Applies a gamma curve to the frame + #Uses a lookup table because that's way faster than actually calculating every pixel + #Immensely improves tracking in low light situations + def applyGamma(self): + img_yuv = cv2.cvtColor(self.frame, cv2.COLOR_BGR2YUV) + #saving a copy of the brightness channel so I can use it to adjust the gamma based on where the face is in that frame + self.brightnessFrame = img_yuv[:,:,0].copy() + #building a lookup table on the fly + lookupTable = np.array(range(256)) + loopupTable = lookupTable/255 + lookupTable = np.power(loopupTable, self.gamma)*255 + img_yuv[:,:,0] = lookupTable[img_yuv[:,:,0]] + #I convert the image to RBG here because it was getting repeatedly converted in the face tracking + self.frame = cv2.cvtColor(img_yuv, cv2.COLOR_YUV2RGB) + return + + + #calculate the ideal gamma based on the brightness of the user's face + #kind of winging it with the math, but it's a lot better than calculating based on the whole frame + def updateGamma(self): + face_center, face_radius = self.faceQueue.get() + + y1 = int(face_center[1] - face_radius[1]) + y2 = int(face_center[1] + face_radius[1]) + x1 = int(face_center[0] - face_radius[0]) + x2 = int(face_center[0] + face_radius[0]) + + #I was doing random sampling of the image brightness + #and that's probably brighter for the whole frame + #but for just the face this is pretty fast + averageBrightness = np.mean(self.brightnessFrame[y1:y2,x1:x2])/256 + self.gamma = math.log(self.targetBrightness, averageBrightness) + return + + +