-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
159 lines (123 loc) · 6.31 KB
/
train.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
import os.path
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression, RidgeClassifier
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.metrics import accuracy_score
import pickle
import cv2
import numpy as np
from mediapipe.python.solutions import drawing_utils as mp_drawing
from mediapipe.python.solutions import holistic as mp_holistic
df = pd.read_csv('../coordinates.csv')
X = df.drop('class', axis=1)
y = df['class']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1111)
pipelines = {
'lr': make_pipeline(StandardScaler(), LogisticRegression()),
'rc': make_pipeline(StandardScaler(), RidgeClassifier()),
'rf': make_pipeline(StandardScaler(), RandomForestClassifier()),
'gb': make_pipeline(StandardScaler(), GradientBoostingClassifier()),
}
fit_models = {}
for algo, pipeline in pipelines.items():
model = pipeline.fit(X_train, y_train)
fit_models[algo] = model
fit_models['rc'].predict(X_test)
for algo, model in fit_models.items():
yhat = model.predict(X_test)
print(algo, accuracy_score(y_test, yhat))
fit_models['rf'].predict(X_test)
with open('body_language.pkl', 'wb') as f:
pickle.dump(fit_models['rf'], f)
cap = cv2.VideoCapture(0)
# Initiate holistic model
with mp_holistic.Holistic(min_detection_confidence=0.5, min_tracking_confidence=0.5) as holistic:
while cap.isOpened():
ret, frame = cap.read()
# Recolor Feed
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image.flags.writeable = False
# Make Detections
results = holistic.process(image)
# print(results.face_landmarks)
# face_landmarks, pose_landmarks, left_hand_landmarks, right_hand_landmarks
# Recolor image back to BGR for rendering
image.flags.writeable = True
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# 1. Draw face landmarks
mp_drawing.draw_landmarks(image, results.face_landmarks, mp_holistic.FACE_CONNECTIONS,
mp_drawing.DrawingSpec(color=(80, 110, 10), thickness=1, circle_radius=1),
mp_drawing.DrawingSpec(color=(80, 256, 121), thickness=1, circle_radius=1)
)
# 2. Right hand
mp_drawing.draw_landmarks(image, results.right_hand_landmarks, mp_holistic.HAND_CONNECTIONS,
mp_drawing.DrawingSpec(color=(80, 22, 10), thickness=2, circle_radius=4),
mp_drawing.DrawingSpec(color=(80, 44, 121), thickness=2, circle_radius=2)
)
# 3. Left Hand
mp_drawing.draw_landmarks(image, results.left_hand_landmarks, mp_holistic.HAND_CONNECTIONS,
mp_drawing.DrawingSpec(color=(121, 22, 76), thickness=2, circle_radius=4),
mp_drawing.DrawingSpec(color=(121, 44, 250), thickness=2, circle_radius=2)
)
# 4. Pose Detections
mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_holistic.POSE_CONNECTIONS,
mp_drawing.DrawingSpec(color=(245, 117, 66), thickness=2, circle_radius=4),
mp_drawing.DrawingSpec(color=(245, 66, 230), thickness=2, circle_radius=2)
)
# Export coordinates
try:
# Extract Pose landmarks
pose = results.pose_landmarks.landmark
pose_row = list(
np.array([[landmark.x, landmark.y, landmark.z, landmark.visibility] for landmark in pose]).flatten())
# Extract Face landmarks
face = results.face_landmarks.landmark
face_row = list(
np.array([[landmark.x, landmark.y, landmark.z, landmark.visibility] for landmark in face]).flatten())
# Concate rows
row = pose_row + face_row
# # Append class name
# row.insert(0, class_name)
# # Export to CSV
# with open('coords.csv', mode='a', newline='') as f:
# csv_writer = csv.writer(f, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
# csv_writer.writerow(row)
# Make Detections
X = pd.DataFrame([row])
body_language_class = model.predict(X)[0]
body_language_prob = model.predict_proba(X)[0]
print(body_language_class, body_language_prob)
# Grab ear coords
coords = tuple(np.multiply(
np.array(
(results.pose_landmarks.landmark[mp_holistic.PoseLandmark.LEFT_EAR].x,
results.pose_landmarks.landmark[mp_holistic.PoseLandmark.LEFT_EAR].y))
, [640, 480]).astype(int))
cv2.rectangle(image,
(coords[0], coords[1] + 5),
(coords[0] + len(body_language_class) * 20, coords[1] - 30),
(245, 117, 16), -1)
cv2.putText(image, body_language_class, coords,
cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
# Get status box
cv2.rectangle(image, (0, 0), (250, 60), (245, 117, 16), -1)
# Display Class
cv2.putText(image, 'CLASS'
, (95, 12), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
cv2.putText(image, body_language_class.split(' ')[0]
, (90, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
# Display Probability
cv2.putText(image, 'PROB'
, (15, 12), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
cv2.putText(image, str(round(body_language_prob[np.argmax(body_language_prob)], 2))
, (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
except:
pass
cv2.imshow('Raw Webcam Feed', image)
if cv2.waitKey(10) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()