-
Notifications
You must be signed in to change notification settings - Fork 0
/
pose-estimation.py
239 lines (189 loc) · 10.1 KB
/
pose-estimation.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
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
import calculateangle
import cv2
import mediapipe as mp
import numpy as np
mp_drawing = mp.solutions.drawing_utils # when visualing out poses
mp_pose = mp.solutions.pose
# for smoothing function
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt
from scipy.signal import savgol_filter
import seaborn as sns
filenames = ["./videos/JJ1.mp4"]
times = []
time = 0
left_elbow_angles = []
left_shoulder_angles = []
left_wrist_angles = []
left_hip_angles = []
left_knee_angles = []
right_elbow_angles = []
right_shoulder_angles = []
right_wrist_angles = []
right_hip_angles = []
right_knee_angles = []
## setup mediapipe instance
images = []
y_values =[]
for filename in filenames:
# clear list
images = []
left_elbow_angles.clear()
times = []
time = 0
time_for_seaborn= []
with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose: #.Pose access pose estimation model, #min_tracking_confidence tracks state
cap = cv2.VideoCapture(filename)
while cap.isOpened():
ret, frame = cap.read() # frame is image from camera
if not ret:
cap.release()
break
frame_width = int(cap.get(3))
# Recolor image
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image.flags.writeable = False # save memory
images.append(image)
# Make detection
results = pose.process(image) # image here is RGB
# Recolor back to BGR
image.flags.writeable = True
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
# Extract landmarks
try:
landmarks = results.pose_landmarks.landmark # hold landamrks. including x,y,z. Use this for calculating angles
# print(landmarks[mp_pose.PoseLandmark.LEFT_FOOT_INDEX])
# Filter out landmarks with low visibility
left_shoulder = [landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].x, landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER.value].y]
left_elbow = [landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].x, landmarks[mp_pose.PoseLandmark.LEFT_ELBOW.value].y]
left_wrist = [landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].x, landmarks[mp_pose.PoseLandmark.LEFT_WRIST.value].y]
right_shoulder = [landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value].x, landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER.value].y]
right_elbow = [landmarks[mp_pose.PoseLandmark.RIGHT_ELBOW.value].x, landmarks[mp_pose.PoseLandmark.RIGHT_ELBOW.value].y]
right_wrist = [landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value].x, landmarks[mp_pose.PoseLandmark.RIGHT_WRIST.value].y]
# shoulder angle
left_hip = [landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].x, landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].y]
right_hip = [landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value].x, landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value].y]
# wrist angle
left_index = [landmarks[mp_pose.PoseLandmark.LEFT_INDEX.value].x, landmarks[mp_pose.PoseLandmark.LEFT_INDEX.value].y]
right_index = [landmarks[mp_pose.PoseLandmark.RIGHT_INDEX.value].x, landmarks[mp_pose.PoseLandmark.RIGHT_INDEX.value].y]
# hip angle
left_knee = [landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value].x, landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value].y]
right_knee = [landmarks[mp_pose.PoseLandmark.RIGHT_KNEE.value].x, landmarks[mp_pose.PoseLandmark.RIGHT_KNEE.value].y]
# knee angle
left_ankle = [landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value].x, landmarks[mp_pose.PoseLandmark.LEFT_ANKLE.value].y]
right_ankle = [landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE.value].x, landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE.value].y]
left_elbow_angle = calculateangle.calculate_angle(left_shoulder, left_elbow, left_wrist)
left_shoulder_angle = calculateangle.calculate_angle(left_hip, left_shoulder, left_elbow)
left_wrist_angle = calculateangle.calculate_angle(left_elbow, left_wrist, left_index)
left_hip_angle = calculateangle.calculate_angle(left_knee, left_hip, left_shoulder)
left_knee_angle = calculateangle.calculate_angle(left_ankle, left_knee, left_hip)
right_elbow_angle = calculateangle.calculate_angle(right_shoulder, right_elbow, right_wrist)
right_shoulder_angle = calculateangle.calculate_angle(right_hip, right_shoulder, right_elbow)
right_wrist_angle = calculateangle.calculate_angle(right_elbow, right_wrist, right_index)
right_hip_angle = calculateangle.calculate_angle(right_knee, right_hip, right_shoulder)
right_knee_angle = calculateangle.calculate_angle(right_ankle, right_knee, right_hip)
# if time % 5 == 0:
left_elbow_angles.append(left_elbow_angle)
left_shoulder_angles.append(left_shoulder_angle)
left_wrist_angles.append(left_wrist_angle)
left_hip_angles.append(left_hip_angle)
left_knee_angles.append(left_knee_angle)
right_elbow_angles.append(right_elbow_angle)
right_shoulder_angles.append(right_shoulder_angle)
right_wrist_angles.append(right_wrist_angle)
right_hip_angles.append(right_hip_angle)
right_knee_angles.append(right_knee_angle)
times.append(time)
time+=1
# Visualize left_elbow position on each frame
cv2.putText(image,
str(left_elbow),
tuple(np.multiply(left_elbow, [640, 480]).astype(int)), # controal [640, 480] to window size
cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,255), 2)
except:
pass
mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS,
mp_drawing.DrawingSpec(color=(245,117,66), thickness=2, circle_radius=2),
mp_drawing.DrawingSpec(color=(245,66,230), thickness=2, circle_radius=2)) # image here is BGR
cv2.imshow('Mediapipe Feed', image)
if cv2.waitKey(10) & 0xFF == ord('q'):
break
# while loop is ended
cap.release()
cv2.destroyAllWindows()
# plot angles vs times graph
print(f'Before sampling: {len(times)}')
plt.plot(times, left_elbow_angles, color='b', label = 'left_elbow')
plt.savefig(f"./output-images/{filename[9:-4]}'s unsmooth-angle-vs-time.jpg")
plt.close()
# sampling
samples = np.random.choice(np.arange(len(times)), size=10, replace=False)
print(samples)
sampled_left_elbow_angles = []
for i, angle in enumerate(left_elbow_angles):
if i not in samples:
print(i)
sampled_left_elbow_angles.append(angle)
sampled_times = np.arange(len(times)-10)
print(sampled_left_elbow_angles)
sampled_left_elbow_angles = np.array(sampled_left_elbow_angles)
print(f'After sampling: {len(sampled_times)}')
plt.plot(sampled_times, sampled_left_elbow_angles, color='b', label = 'sampled_left_elbow')
plt.savefig(f"./output-images/{filename[9:-4]}'s sampled-unsmooth-angle-vs-time.jpg")
##### smooth a curve #####
# cubic_interpolation_model = interp1d(times, left_elbow_angles, kind = "cubic")
cubic_interpolation_model = interp1d(sampled_times, sampled_left_elbow_angles, kind = "cubic")
sampled_times = np.array(sampled_times)
Y_ = cubic_interpolation_model(sampled_times)
# apply smoothing filter
Y_ = savgol_filter(Y_, len(sampled_times), 50)
# for margin graph
y = -Y_
y_values += list(y)
##### smooth a curve #####
## plot each angle
# plot left angles
# plt.plot(times,Y_, color='r', label='smooth')
plt.plot(sampled_times,Y_, color='r', label='smooth')
# plt.plot(times, left_shoulder_angles, color='r', label='left_shoulder')
# plt.plot(times, left_wrist_angles, color='g', label = 'left wrist')
# plt.plot(times, left_hip_angles, color='y', label = 'left hip')
# plt.plot(times, left_knee_angles, color='m', label = 'left knee')
# plot right angles
# plt.plot(times, right_elbow_angles, color='#87CEEB', label = 'right_elbow')
# plt.plot(times, right_shoulder_angles, color='#FFC0CB', label='right_shoulder')
# plt.plot(times, right_wrist_angles, color='#90EE90', label = 'right wrist')
# plt.plot(times, right_hip_angles, color='#FFF01F', label = 'right hip')
# plt.plot(times, right_knee_angles, color='#A020F0', label = 'right knee')
plt.xlabel('time')
plt.ylabel('angle')
plt.legend(title=f"{filename[9:-4]}'s angle vs time")
plt.savefig(f"./output-images/{filename[9:-4]}'s sampled smoothed angle-vs-time.jpg")
plt.show(block=False)
plt.close()
# visualize margin
sns.set()
print(f'y_values length: {len(y_values)}')
y_values = np.array(y_values).reshape(2, -1)
y_values = np.mean(y_values, axis=0)
print(f'y_values.shape: {y_values.shape}')
x_axis = np.arange(len(y_values))
plt.plot(x_axis, y_values, 'b-', label='y_value')
plt.fill_between(x_axis, y_values-10, y_values+10, color = 'b', alpha=0.2)
plt.legend(title='margin')
plt.ioff()
# plt.savefig
plt.show(block=False)
plt.savefig("./output-images/margin.jpg")
# def show_margin(y_values):
# sns.set()
# print(f'y_values length: {len(y_values)}')
# y_values = np.array(y_values).reshape(2, -1)
# y_values = np.mean(y_values, axis=0)
# print(f'y_values.shape: {y_values.shape}')
# x_axis = np.arange(len(y_values))
# plt.plot(x_axis, y_values, 'b-', label='y_value')
# plt.fill_between(x_axis, y_values-10, y_values+10, color = 'b', alpha=0.2)
# plt.legend(title='test')
# plt.show()
# show_margin(y_values)