-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathutils.py
173 lines (140 loc) · 8.28 KB
/
utils.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
166
167
168
## UNCOMMENT BOTTOM LINES AND RUN THIS ONCE
import mediapipe as mp
import cv2
import pandas as pd
import os
from sklearn.metrics import confusion_matrix, classification_report
import numpy as np
mp_pose = mp.solutions.pose
pose = mp_pose.Pose()
points = mp_pose.PoseLandmark # Landmarks
mp_drawing = mp.solutions.drawing_utils # For drawing keypoints
def calculate_angle(landmark1, landmark2, landmark3):
x1, y1, _ = landmark1.x, landmark1.y, landmark1.z
x2, y2, _ = landmark2.x, landmark2.y, landmark2.z
x3, y3, _ = landmark3.x, landmark3.y, landmark3.z
angle = np.degrees(np.arctan2(y3 - y2, x3 - x2) - np.arctan2(y1 - y2, x1 - x2))
# Check if the angle is less than zero.
if angle < 0:
# Add 360 to the found angle.
angle += 360
return angle
def extract_pose_angles(results):
angles = []
if results.pose_landmarks is not None:
landmarks = results.pose_landmarks.landmark
# Get the angle between the left elbow, wrist and left index points.
left_wrist_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value],
landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value],
landmarks[mp_pose.PoseLandmark.LEFT_INDEX.value])
angles.append(left_wrist_angle)
# Get the angle between the right elbow, wrist and left index points.
right_wrist_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.RIGHT_ELBOW.value],
landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value],
landmarks[mp_pose.PoseLandmark.RIGHT_INDEX.value])
angles.append(right_wrist_angle)
# Get the angle between the left shoulder, elbow and wrist points.
left_elbow_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value],
landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value],
landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value])
angles.append(left_elbow_angle)
# Get the angle between the right shoulder, elbow and wrist points.
right_elbow_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value],
landmarks[mp_pose.PoseLandmark.RIGHT_ELBOW.value],
landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value])
angles.append(right_elbow_angle)
# Get the angle between the left elbow, shoulder and hip points.
left_shoulder_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value],
landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value],
landmarks[mp_pose.PoseLandmark.LEFT_HIP.value])
angles.append(left_shoulder_angle)
# Get the angle between the right hip, shoulder and elbow points.
right_shoulder_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value],
landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value],
landmarks[mp_pose.PoseLandmark.RIGHT_ELBOW.value])
angles.append(right_shoulder_angle)
# Get the angle between the left hip, knee and ankle points.
left_knee_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.LEFT_HIP.value],
landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value],
landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value])
angles.append(left_knee_angle)
# Get the angle between the right hip, knee and ankle points
right_knee_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value],
landmarks[mp_pose.PoseLandmark.RIGHT_KNEE.value],
landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE.value])
angles.append(right_knee_angle)
# Get the angle between the left hip, ankle and LEFT_FOOT_INDEX points.
left_ankle_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.LEFT_HIP.value],
landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value],
landmarks[mp_pose.PoseLandmark.LEFT_FOOT_INDEX.value])
angles.append(left_ankle_angle)
# Get the angle between the right hip, ankle and RIGHT_FOOT_INDEX points
right_ankle_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value],
landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE.value],
landmarks[mp_pose.PoseLandmark.RIGHT_FOOT_INDEX.value])
angles.append(right_ankle_angle)
# Get the angle between the left knee, hip and right hip points.
left_hip_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value],
landmarks[mp_pose.PoseLandmark.LEFT_HIP.value],
landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value])
angles.append(left_hip_angle)
# Get the angle between the left hip, right hip and right kneee points
right_hip_angle = calculate_angle(landmarks[mp_pose.PoseLandmark.LEFT_HIP.value],
landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value],
landmarks[mp_pose.PoseLandmark.RIGHT_KNEE.value])
angles.append(right_hip_angle)
return angles
# Predict the name of the poses in the image
def predict(img, model, show=False):
img = cv2.imread(img)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
#Resize the image to 50% of the original size
scale_factor = 0.5
img = cv2.resize(img, None, fx=scale_factor, fy=scale_factor)
results = pose.process(img)
if results.pose_landmarks:
list_angles = []
list_angles = extract_pose_angles(results)
y = model.predict([list_angles])
if show:
mp_drawing.draw_landmarks(img, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
cv2.putText(img, str(y[0]), (50,50), cv2.FONT_HERSHEY_SIMPLEX,1,(215,215,0),3)
cv2.imshow("image", img)
cv2.waitKey(0)
def predict_video(model, video="0", show=False):
cap = cv2.VideoCapture(video)
while cap.isOpened():
angles = []
success, img = cap.read()
if not success:
print("Ignoring empty camera frame.")
continue
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
results = pose.process(img)
if results.pose_landmarks:
list_angles = []
list_angles = extract_pose_angles(results)
y = model.predict([list_angles])
name = str(y[0])
if show:
mp_drawing.draw_landmarks(img, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
(w, h), _ = cv2.getTextSize(name, cv2.FONT_HERSHEY_SIMPLEX, 1, 1)
cv2.rectangle(img, (40, 40), (40+w, 60), (255, 255, 255), cv2.FILLED)
cv2.putText(img, name, (40, 60), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 0), 3)
cv2.imshow("Video", img)
if cv2.waitKey(5) & 0xFF == 27:
break
cap.release()
# Use this to evaluate any dataset you've built
def evaluate(data_test, model, show=False):
target = data_test.loc[:, "target"] # list of labels
target = target.values.tolist()
predictions = []
for i in range(len(data_test)):
tmp = data_test.iloc[i, 0:len(data_test.columns) - 1]
tmp = tmp.values.tolist()
predictions.append(model.predict([tmp])[0])
if show:
print(confusion_matrix(predictions, target), '\n')
print(classification_report(predictions, target))
return predictions