diff --git a/facedetection.py b/facedetection.py index db84b2a..7663fba 100644 --- a/facedetection.py +++ b/facedetection.py @@ -5,15 +5,18 @@ def detect_faces(frame, model, detection_threshold ): image = resizeImage(frame) - outputs, _ = model.faceDetection.run([], {'input': image}) + outputs, maxpool = model.faceDetection.run([], {'input': image}) outputs = outputs[0] + maxpool = maxpool[0] + mask = outputs[0] == maxpool[0] + mask2 = outputs[1] > 0.2 + mask = mask.astype(int) * mask2.astype(int) + outputs[0] = outputs[0] * mask.astype(int) faceLocation = np.argmax(outputs[0].flatten()) x = faceLocation % 56 y = faceLocation // 56 - if outputs[0, y, x] < detection_threshold: 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) diff --git a/tracker.py b/tracker.py index b4de59d..9b1bf77 100644 --- a/tracker.py +++ b/tracker.py @@ -48,7 +48,7 @@ class Models(): 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, featureType, threads, model_type=3, detection_threshold=0.6, threshold=0.6, silent=False): + def __init__(self, width, height, featureType, threads, model_type=3, detection_threshold=0.5, threshold=0.6, silent=False): self.detection_threshold = detection_threshold self.EyeTracker = eyes.EyeTracker()