-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
122 lines (106 loc) · 5 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import cv2
import numpy as np
import dlib
import pandas as pd
from imutils import face_utils
import os
from object_detection.utils import label_map_util
from object_detection.builders import model_builder
from object_detection.utils import config_util
import tensorflow as tf
from tensorflow.keras.models import load_model
detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor('track dependencies\shape_predictor_68_face_landmarks.dat')
CONFIG_PATH = 'pupil track dependencies/pipeline.config'
CHECKPOINT_PATH = 'pupil track dependencies/'
WEIGHT_PATH = "4layer/Weights-9788--0.92333--0.96371.hdf5"
configs = config_util.get_configs_from_pipeline_file(CONFIG_PATH)
detection_model = model_builder.build(model_config=configs['model'], is_training=False)
ckpt = tf.compat.v2.train.Checkpoint(model=detection_model)
ckpt.restore(os.path.join(CHECKPOINT_PATH, 'ckpt-7')).expect_partial()
category_index = label_map_util.create_category_index_from_labelmap('pupil track dependencies/label_map.pbtxt')
model = load_model('4layer/Model.h5')
model.load_weights(WEIGHT_PATH)
@tf.function
def detect_pupil(image):
image, shapes = detection_model.preprocess(image)
prediction_dict = detection_model.predict(image, shapes)
detections = detection_model.postprocess(prediction_dict, shapes)
return detections
def main_pupil(frame):
image_np = np.array(frame)
input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32)
detections = detect_pupil(input_tensor)
num_detections = int(detections.pop('num_detections'))
detections = {key: value[0, :num_detections].numpy()
for key, value in detections.items()}
box = np.squeeze(detections['detection_boxes'])
j = []
for i in range(3):
ymin = (int(box[i, 0] * 480))
xmin = (int(box[i, 1] * 640))
ymax = (int(box[i, 2] * 480))
xmax = (int(box[i, 3] * 640))
z = (xmin, ymin, xmax, ymax)
j.append(z)
return j[0][0], j[0][1], j[0][2], j[0][3], j[1][0], j[1][1], j[1][2], j[1][3], j[2][0], j[2][1], j[2][2], j[2][3]
def get_face_points(face):
(x, y, w, h) = face_utils.rect_to_bb(face)
return x,y,w,h
def get_head_pose(img, facial_landmarks):
hp = [30,8,36,45,59,55]
size = img.shape
image_points = np.array([(facial_landmarks.part(point).x, facial_landmarks.part(point).y) for point in hp], dtype="double")
model_points = np.array([(0.0, 0.0, 0.0),(0.0, -330.0, -65.0),(-225.0, 170.0, -135.0),(225.0, 170.0, -135.0),(-150.0, -150.0, -125.0),(150.0, -150.0, -125.0)])
focal_length = size[1]
center = (size[1] / 2, size[0] / 2)
camera_matrix = np.array([[focal_length, 0, center[0]],[0, focal_length, center[1]],[0, 0, 1]], dtype="double")
dist_coeffs = np.zeros((4, 1))
(success, rotation_vector, translation_vector) = cv2.solvePnP(model_points, image_points, camera_matrix, dist_coeffs)
(nose_end_point2D, jacobian) = cv2.projectPoints(np.array([(0.0, 0.0, 1000.0)]), rotation_vector,
translation_vector, camera_matrix, dist_coeffs)
p1 = (int(image_points[0][0]), int(image_points[0][1]))
p2 = (int(nose_end_point2D[0][0][0]), int(nose_end_point2D[0][0][1]))
return p1[0],p1[1],p2[0],p2[1]
def get_eye_points(img, facial_landmarks):
lp = [36, 37, 38, 39, 40, 41]
rp = [42, 43, 44, 45, 46, 47]
mask = np.zeros((480, 640), np.uint8)
lr = np.array([(facial_landmarks.part(point).x, facial_landmarks.part(point).y) for point in lp])
rr = np.array([(facial_landmarks.part(point).x, facial_landmarks.part(point).y) for point in rp])
mask1 = cv2.fillPoly(mask, [lr], 255)
mask1 = cv2.fillPoly(mask1, [rr], 255)
eye = cv2.bitwise_and(img, img, mask=mask1)
leftEyeCenter = lr.mean(axis=0).astype("int")
rightEyeCenter = rr.mean(axis=0).astype("int")
dY = rightEyeCenter[1] - leftEyeCenter[1]
dX = rightEyeCenter[0] - leftEyeCenter[0]
angle = np.degrees(np.arctan2(dY, dX))
angle = round(angle, 2)
return eye, angle
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
while True:
_, frame = cap.read()
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
faces = detector(gray)
key = cv2.waitKey(1)
# cv2.imshow('ouput', frame)
if (len(faces) >= 1):
for face in faces:
landmarks = predictor(gray, face)
CC, angle=get_eye_points(frame, landmarks)
x,y,w,he = get_face_points(face)
x0,y0,x1,y1 = get_head_pose(gray, landmarks)
a,b,c,d,e,f,g,h,i,j,k,l = main_pupil(CC)
m = np.array([a,b,c,d,e,f,g,h,i,j,k,l,angle, x,y,w,he,x0,y0,x1,y1])
m = tf.convert_to_tensor(np.expand_dims(m, 0), dtype=tf.float32)
n = model.predict(m)
if (n > 0.5):
print('Yes')
else:
print('No')
else:
continue
if cv2.waitKey(1) & 0xFF == ord('q'):
cap.release()
break