copilot/OpenSeeFace
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

moved computer vision models to one dedicated class

Browse source
This commit is contained in:
Cowwan 2024-02-06 20:24:37 -05:00
parent 797e5c2d56
commit 14744143bf
6 changed files with 64 additions and 97 deletions
Download patch file Download diff file Expand all files Collapse all files

14
eyes.py
View file

@ -149,17 +149,7 @@ class EyeTracker():
self.faceCenter = None
self.faceRadius = None
model_base_path = os.path.join(os.path.dirname(__file__), os.path.join("models"))
providersList = onnxruntime.capi._pybind_state.get_available_providers()
options = onnxruntime.SessionOptions()
options.inter_op_num_threads = 1
options.intra_op_num_threads = 1
options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
options.log_severity_level = 3
self.gaze_model = onnxruntime.InferenceSession(os.path.join(model_base_path, "mnv3_gaze32_split_opt.onnx"), sess_options=options, providers=providersList)
def get_eye_state(self,frame, lms):
def get_eye_state(self, models, frame, lms):
lms = self.extract_face(frame, np.array(lms)[:,0:2][:,::-1])
@ -177,7 +167,7 @@ class EyeTracker():
return [[1.,0.,0.,0.],[1.,0.,0.,0.]] #Early exit if one of the eyes doesn't have data
both_eyes = np.concatenate((self.rightEye.image, self.leftEye.image))
self.rightEye.results, self.leftEye.results = self.gaze_model.run([], {"input": both_eyes})[0]
self.rightEye.results, self.leftEye.results = models.gazeTracker.run([], {"input": both_eyes})[0]
self.rightEye.calculateEye()
self.leftEye.calculateEye()

4
face.py
View file

@ -4,10 +4,8 @@ import featureExtractor
import emilianaFeatureExtractor
class FaceInfo():
def __init__(self, id, tracker, featureType):
def __init__(self, featureType):
self.featureType = featureType
self.id = id
self.tracker = tracker
self.face_3d = np.array([
[ 0.4551769692672 , 0.300895790030204, -0.764429433974752],
[ 0.448998827123556, 0.166995837790733, -0.765143004071253],

38
facedetection.py
View file

@ -5,34 +5,9 @@ import onnxruntime
import time
#kinda split on how to handle this one
#It's the cause of all of my slowest frames
#but it's not run often
#and I don't know how to make it faster
class FaceDetectionModel():
def __init__(self):
model_base_path = os.path.join(os.path.dirname(__file__), os.path.join("models"))
providersList = onnxruntime.capi._pybind_state.get_available_providers()
options = onnxruntime.SessionOptions()
options.inter_op_num_threads = 1
options.intra_op_num_threads = 1
options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
options.log_severity_level = 3
self.detection = onnxruntime.InferenceSession(os.path.join(model_base_path, "mnv3_detection_opt.onnx"), sess_options=options, providers=providersList)
def run(self, image):
outputs, maxpool = self.detection.run([], {'input': image})
return outputs[0]
class FaceDetector():
def __init__(self, detection_threshold=0.6):
self.model = FaceDetectionModel()
self.mean = np.float32(np.array([-2.1179, -2.0357, -1.8044]))
self.std = np.float32(np.array([0.0171, 0.0175, 0.0174]))
@ -40,21 +15,24 @@ class FaceDetector():
self.targetimageSize = [224, 224]
self.image = None
def detect_faces(self, frame):
def detect_faces(self, frame, model):
self.resizeImage(frame)
outputs = self.model.run(self.image)
outputs, _ = model.faceDetection.run([], {'input': self.image})
outputs = outputs[0]
faceLocation = np.argmax(outputs[0].flatten())
x = faceLocation % 56
y = faceLocation // 56
if outputs[0, y, x] < self.detection_threshold:
return []
return None
r = outputs[1, y, x] * 112.
results= (((x * 4) - r, (y * 4) - r, r*2,r*2))
results = np.array(results).astype(np.float32)
results[ [0,2]] *= frame.shape[1] / 224.
results[ [1,3]] *= frame.shape[0] / 224.
results[[0,2]] *= frame.shape[1] / 224.
results[[1,3]] *= frame.shape[0] / 224.
return results
def resizeImage(self, frame):

37
landmarks.py
View file

@ -28,37 +28,13 @@ def matrix_to_quaternion(m):
#not because I didn't try to improve things
#but because I caused problems when I did
class Landmarks():
def __init__(self, width, height, model_type, threshold):
#removed all the tracking levels below 0
#because they were so feature bare that they just didn't seem usable
#I beleive the selection should be equivalent to VTS, just 0 indexed
self.model_type = model_type
self.models = [
"lm_model0_opt.onnx",
"lm_model1_opt.onnx",
"lm_model2_opt.onnx",
"lm_model3_opt.onnx",
"lm_model4_opt.onnx"]
model = self.models[self.model_type]
model_base_path = os.path.join(os.path.dirname(__file__), os.path.join("models"))
providersList = onnxruntime.capi._pybind_state.get_available_providers()
options = onnxruntime.SessionOptions()
options.inter_op_num_threads = 1
#the only num_threads that seems to matter
#Doesn't seem to scale well beyond 2
options.intra_op_num_threads = 1
options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
options.log_severity_level = 3
self.session = onnxruntime.InferenceSession(os.path.join(model_base_path, model), sess_options=options, providers=providersList)
self.input_name = self.session.get_inputs()[0].name
def __init__(self, width, height, threshold):
self.camera = np.array([[width, 0, width/2], [0, width, height/2], [0, 0, 1]], np.float32)
self.inverse_camera = np.linalg.inv(self.camera)
self.width = width
self.height = height
self.confidence = 0
self.face_3d = np.array([
[ 0.4551769692672 , 0.300895790030204, -0.764429433974752],
@ -138,10 +114,11 @@ class Landmarks():
self.fail_count = 0
def run(self, crop, crop_info):
output = self.session.run([], {self.input_name: crop})[0][0]
conf, lms = self.landmarks(output, crop_info)
return (conf, lms)
def run(self, models, crop, crop_info):
output = models.landmarks.run([], {"input": crop})[0][0]
self.confidence, lms = self.landmarks(output, crop_info)
return ( lms)
def landmarks(self, tensor, crop_info):
crop_x1, crop_y1, scale_x, scale_y = crop_info

66
tracker.py
View file

@ -1,12 +1,12 @@
import os
import numpy as np
import onnxruntime
import cv2
import time
import facedetection
import eyes
import landmarks
import face
import featureExtractor
import emilianaFeatureExtractor
#this file is so much smaller than it used to be, lol
@ -21,38 +21,63 @@ def clamp_to_im(pt, w, h): #8 times per frame, but that only accounts for 0.005m
y = h-1
return (int(x),int(y+1))
class Models():
def __init__(self, model_type=3):
self.model_type = model_type
self.models = [
"lm_model0_opt.onnx",
"lm_model1_opt.onnx",
"lm_model2_opt.onnx",
"lm_model3_opt.onnx",
"lm_model4_opt.onnx"]
model = self.models[self.model_type]
model_base_path = os.path.join(os.path.dirname(__file__), os.path.join("models"))
providersList = onnxruntime.capi._pybind_state.get_available_providers()
options = onnxruntime.SessionOptions()
options.inter_op_num_threads = 1
options.intra_op_num_threads = 1
options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
options.log_severity_level = 3
self.landmarks = onnxruntime.InferenceSession(os.path.join(model_base_path, model), sess_options=options, providers=providersList)
self.faceDetection = onnxruntime.InferenceSession(os.path.join(model_base_path, "mnv3_detection_opt.onnx"), sess_options=options, providers=providersList)
self.gazeTracker = onnxruntime.InferenceSession(os.path.join(model_base_path, "mnv3_gaze32_split_opt.onnx"), sess_options=options, providers=providersList)
class Tracker():
def __init__(self, width, height, messageQueue, featureType, model_type=3, detection_threshold=0.6, threshold=0.6, silent=False):
self.FaceDetector = facedetection.FaceDetector(detection_threshold = detection_threshold)
self.EyeTracker = eyes.EyeTracker()
self.Landmarks = landmarks.Landmarks(width, height, model_type, threshold)
if featureType == 0:
self.FeatureExtractor = featureExtractor.FeatureExtractor()
else:
self.FeatureExtractor = emilianaFeatureExtractor.FeatureExtractor()
self.Landmarks = landmarks.Landmarks(width, height, threshold)
self.model = Models(model_type = model_type)
# Image normalization constants
self.mean = np.float32(np.array([-2.1179, -2.0357, -1.8044]))
self.std = np.float32(np.array([0.0171, 0.0175, 0.0174]))
self.targetimageSize = [224, 224]
self.width = width
self.height = height
self.threshold = threshold
self.face = None
self.face_info = [face.FaceInfo(0, self, featureType)]
self.face_info = face.FaceInfo(featureType)
self.messageQueue = messageQueue
#scales cropped frame before sending to landmarks
def preprocess(self, im):
im = cv2.resize(im, (224, 224), interpolation=cv2.INTER_CUBIC) * self.std + self.mean
im = cv2.resize(im ,self.targetimageSize , interpolation=cv2.INTER_CUBIC) * self.std + self.mean
im = np.expand_dims(im, 0)
return np.transpose(im, (0,3,1,2))
def cropFace(self,face, im):
def cropFace(self, im):
duration_pp = 0.0
x,y,w,h = face
x,y,w,h = self.face
wint = int(w * 0.1)
hint = int(h * 0.125)
@ -94,30 +119,29 @@ class Tracker():
if self.face is None:
start_fd = time.perf_counter()
self.face = self.FaceDetector.detect_faces(frame)
self.face = self.FaceDetector.detect_faces(frame, self.model)
duration_fd = 1000 * (time.perf_counter() - start_fd)
if self.face is None:
return self.early_exit("No faces found", start)
crop, crop_info, duration_pp = self.cropFace(self.face, frame)
crop, crop_info, duration_pp = self.cropFace(frame)
#Early exit if crop fails, If the crop fails there's nothing to track
if crop is None:
return self.early_exit("No valid crops", start)
start_model = time.perf_counter()
conf, lms = self.Landmarks.run(crop , crop_info )
lms = self.Landmarks.run(self.model, crop, crop_info)
#Early exit if below confidence threshold
if conf < self.threshold:
if self.Landmarks.confidence < self.threshold:
return self.early_exit("Confidence below threshold", start)
eye_state = self.EyeTracker.get_eye_state(frame, lms)
eye_state = self.EyeTracker.get_eye_state(self.model, frame, lms)
self.face_info[0].update((conf, (lms, eye_state)), np.array(lms)[:, 0:2].mean(0))
self.face_info.update((self.Landmarks.confidence, (lms, eye_state)), np.array(lms)[:, 0:2].mean(0))
duration_model = 1000 * (time.perf_counter() - start_model)
start_pnp = time.perf_counter()
face_info = self.face_info[0]
face_info = self.face_info
if face_info.alive:
face_info.success, face_info.quaternion, face_info.euler, face_info.pnp_error, face_info.pts_3d, face_info.lms = self.Landmarks.estimate_depth(face_info)

2
vts.py
View file

@ -30,7 +30,7 @@ class VTS():
print(f"Confidence: {face.conf:.4f} / 3D fitting error: {face.pnp_error:.4f} / Eyes: {left_state}, {right_state}")
now = time.time()
packet.extend(bytearray(struct.pack("d", now)))
packet.extend(bytearray(struct.pack("i", face.id)))
packet.extend(bytearray(struct.pack("i", 0)))
packet.extend(bytearray(struct.pack("f", self.width)))
packet.extend(bytearray(struct.pack("f", self.height)))
packet.extend(bytearray(struct.pack("f", face.eye_blink[0])))