-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
165 lines (125 loc) · 5.96 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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
# Image Processing
import cv2 # For computer vision and image processing
import numpy as np # Support for large, multi-dimensional arrays and matrices of numerical data and collection of mathematical functions
import mediapipe as mp # For face mesh detection
import pandas as pd # For creating a table of head movements
import matplotlib.pyplot as plt # For plotting graphs
from head_detection import detect_head # Import the head detection function
def start_video_capture():
cap = cv2.VideoCapture(0) # Use the first camera device it finds
# Initialize Mediapipe Face Mesh
mp_face_mesh = mp.solutions.face_mesh
face_mesh = mp_face_mesh.FaceMesh(min_detection_confidence=0.5, min_tracking_confidence=0.5)
# List to store head movement data
head_movements = []
while True:
ret, frame = cap.read() # Capture frame-by-frame
if not ret:
print("Error: Can't receive frame")
break
# Flip the image horizontally for a later selfie-view display
# Also convert the color space from BGR to RGB
image = cv2.cvtColor(cv2.flip(frame, 1), cv2.COLOR_BGR2RGB)
# To improve performance
image.flags.writeable = False
# Get the result
results = face_mesh.process(image)
# To improve performance
image.flags.writeable = True
# Convert the color space from RGB to BGR
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
img_h, img_w, img_c = image.shape
face_3d = []
face_2d = []
if results.multi_face_landmarks:
for face_landmarks in results.multi_face_landmarks:
# Calculate the bounding box coordinates
min_x = min(int(lm.x * img_w) for lm in face_landmarks.landmark)
min_y = min(int(lm.y * img_h) for lm in face_landmarks.landmark)
max_x = max(int(lm.x * img_w) for lm in face_landmarks.landmark)
max_y = max(int(lm.y * img_h) for lm in face_landmarks.landmark)
# Draw the bounding box
cv2.rectangle(image, (min_x, min_y), (max_x, max_y), (0, 255, 0), 2)
for idx, lm in enumerate(face_landmarks.landmark):
if idx == 33 or idx == 263 or idx == 1 or idx == 61 or idx == 291 or idx == 199:
if idx == 1:
nose_2d = (lm.x * img_w, lm.y * img_h)
nose_3d = (lm.x * img_w, lm.y * img_h, lm.z * 8000)
x, y = int(lm.x * img_w), int(lm.y * img_h)
# Get the 2D Coordinates
face_2d.append([x, y])
# Get the 3D Coordinates
face_3d.append([x, y, lm.z])
# Convert it to the NumPy array
face_2d = np.array(face_2d, dtype=np.float64)
# Convert it to the NumPy array
face_3d = np.array(face_3d, dtype=np.float64)
# The camera matrix
focal_length = 1 * img_w
cam_matrix = np.array([ [focal_length, 0, img_h / 2],
[0, focal_length, img_w / 2],
[0, 0, 1]])
# The Distance Matrix
dist_matrix = np.zeros((4, 1), dtype=np.float64)
# Solve PnP
success, rot_vec, trans_vec = cv2.solvePnP(face_3d, face_2d, cam_matrix, dist_matrix)
# Get rotational matrix
rmat, jac = cv2.Rodrigues(rot_vec)
# Get angles
angles, mtxR, mtxQ, Qx, Qy, Qz = cv2.RQDecomp3x3(rmat)
# Get the y rotation degree
x = angles[0] * 360
y = angles[1] * 360
# See where the user's head tilting
if y < -10:
text = "Looking Left"
elif y > 10:
text = "Looking Right"
elif x < -10:
text = "Looking Down"
else:
text = "Forward"
# Append head movement data
head_movements.append({
"frame": len(head_movements),
"x_angle": x,
"y_angle": y,
"direction": text
})
# Display the nose direction
nose_3d_projection, jacobian = cv2.projectPoints(nose_3d, rot_vec, trans_vec, cam_matrix, dist_matrix)
p1 = (int(nose_2d[0]), int(nose_2d[1]))
p2 = (int(nose_3d_projection[0][0][0]), int(nose_3d_projection[0][0][1]))
cv2.line(image, p1, p2, (255, 0, 0), 2)
# Add the text on the image
cv2.putText(image, text, (20, 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
cv2.imshow('Frame', image) # Display the frame
key = cv2.waitKey(1) # Wait for a key press, 1ms delay
if key == 27:
break
# Release the camera and destroy all windows after the loop has finished
cap.release()
cv2.destroyAllWindows()
# Create a DataFrame from the head movements data
df = pd.DataFrame(head_movements)
# Plot the head movement directions over time
plt.figure(figsize=(14, 6))
# Plot for head directions
plt.subplot(1, 2, 1)
plt.plot(df['frame'], df['direction'], marker='o')
plt.title('Head Movement Direction Over Time')
plt.xlabel('Frame')
plt.ylabel('Direction')
plt.xticks(rotation=45)
# Plot for head angles
plt.subplot(1, 2, 2)
plt.plot(df['frame'], df['x_angle'], label='X Angle', marker='o')
plt.plot(df['frame'], df['y_angle'], label='Y Angle', marker='x')
plt.title('Head Rotation Angles Over Time')
plt.xlabel('Frame')
plt.ylabel('Angle (degrees)')
plt.legend()
plt.tight_layout()
plt.show()
# Call the function to start the video capture
start_video_capture()