-
Notifications
You must be signed in to change notification settings - Fork 6
/
pose_estimation.py
165 lines (141 loc) · 6.78 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
#==================================================================================
# POSE ESTIMATION
#----------------------------------------------------------------------------------
# Input: Video x 2, Output: JSON
# Given a front/back view and side view video of someone
# walking, this will generate a json, describing the pose
# via key-points in graph form, throughout every frame.
#==================================================================================
#==================================================================================
# Imports
#==================================================================================
from __future__ import division
from gluoncv.data.transforms.pose import detector_to_alpha_pose, heatmap_to_coord_alpha_pose
from gluoncv.model_zoo import get_model
from tqdm import tqdm
import gluoncv as gcv
import mxnet as mx
import time, cv2
import json
import glob
import os
#==================================================================================
# AI Detectors
# Person detector; YOLO via GPU.
# Pose Estimator; AlphaPose via CPU (no GPU support).
#==================================================================================
ctx = mx.gpu(0)
detector = get_model('yolo3_mobilenet1.0_coco', pretrained=True, ctx=ctx)
detector.reset_class(classes=['person'], reuse_weights={'person': 'person'})
estimator = get_model('alpha_pose_resnet101_v1b_coco', pretrained='de56b871')
detector.hybridize()
estimator.hybridize()
#==================================================================================
# Methods
#==================================================================================
# Returns a keypoints with respect to the current frame's pose
def curr_pose(img, coords, confidence, scores, keypoint_thresh=0.2):
i = scores.argmax() # gets index of most confident bbox estimation
pts = coords[i] # coords of most confident pose in frame
pose_data = []
for j in range(0, len(pts)):
x = -1
y = -1
if(confidence[i][j] > keypoint_thresh):
x = int(pts[j][0])
y = int(img.shape[0] - pts[j][1])
pose_data.append([x,y])
return pose_data
# Given one video, returns list of pose information in preparation for json file
def video_to_listPose(vid):
cap = cv2.VideoCapture(vid) # load video
if (cap.isOpened() == False): # Check if camera opened successfully
print("Error opening video stream or file")
return
frame_count = 0
pose_data_vid = []
dimensions = (0, 0)
frame_length = (int)(cap.get(cv2.CAP_PROP_FRAME_COUNT))
pbar = tqdm(total=frame_length, ncols=1, desc='.')
pbar.ncols = 100
# Iterate through every frame in video
while(cap.isOpened()):
ret, frame = cap.read() # read current frame
if (frame is None):
break # If current frame doesn't exist, finished iterating through frames
frame = mx.nd.array(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).astype('uint8') # mxnet readable
# Person detection
x, frame = gcv.data.transforms.presets.yolo.transform_test(frame) # short=406, max_size=1024
class_IDs, scores, bounding_boxs = detector(x.as_in_context(ctx))
# Pose estimation
pose_input, upscale_bbox = detector_to_alpha_pose(frame, class_IDs, scores, bounding_boxs,
output_shape=(320, 256))
# Gets current pose keypoints
if (upscale_bbox is None): # Caters for no person detection
pbar.set_description_str('Skipping ')
pose_data_curr = [[-1, -1] for j in range(0, 17)]
else: # Caters for person detection
pbar.set_description_str('Processing')
predicted_heatmap = estimator(pose_input)
pred_coords, confidence = heatmap_to_coord_alpha_pose(predicted_heatmap, upscale_bbox)
scores = scores.asnumpy()
confidence = confidence.asnumpy()
pred_coords = pred_coords.asnumpy()
# Preparing for json
pose_data_curr = curr_pose(frame, pred_coords, confidence, scores, keypoint_thresh=0.2)
pose_data_vid.append(pose_data_curr)
if (frame_count == 0):
dimensions = [frame.shape[1], frame.shape[0]]
frame_count += 1
pbar.update(1)
cap.release()
pbar.close()
return dimensions, pose_data_vid
# Given two videos it will output the json describing all poses in both videos
def videos_to_jsonPose(vidSide, vidFront, partId, capId):
dimensions_side, pose_vid_side = video_to_listPose(vidSide)
dimensions_front, pose_vid_front = video_to_listPose(vidFront)
if(dimensions_side != dimensions_front):
print('Warning: side video', dimensions_side, 'and front video', dimensions_front, 'of different dimensions' )
if (len(pose_vid_side) != len(pose_vid_front)):
print('Warning: side video', len(pose_vid_side), 'and front video', len(pose_vid_front), 'of different frame counts')
jsonPose_dict = {
'partId': partId,
'capId': capId,
'dimS': dimensions_side,
'lenS' : len(pose_vid_side),
'dimF' : dimensions_front,
'lenF' : len(pose_vid_front),
'dataS' : pose_vid_side,
'dataF' : pose_vid_front
}
return jsonPose_dict
def estimate_poses(path, writeFile):
jsonPose_list = []
fs_pair = []
i = 0
for filename in glob.glob(os.path.join(path, '*.avi')):
fs_pair.append(filename)
if (i % 2):
print('Capture pair', '('+ str(int((i+1)/2)) +'/6)', ':', '\"'+fs_pair[1]+'\"', ',', '\"'+fs_pair[0]+'\"')
capId = fs_pair[0].split('-')[1]
partId = fs_pair[0].split('-')[0].split('\\')[2]
jsonPose_dict = videos_to_jsonPose(fs_pair[1], fs_pair[0], partId, capId)
jsonPose_list.append(jsonPose_dict)
fs_pair.clear()
i += 1
with open(writeFile, 'w') as outfile:
json.dump(jsonPose_list, outfile, separators=(',', ':'))
#==================================================================================
# Main
#==================================================================================
def main():
for i in range(1, 22):
if (len(str(i)) < 2): i = '0' + str(i)
path = '..\\Part' + str(i) + '\\'
writeFile = path + 'Part' + str(i) + '_pose.json'
start_time = time.time()
estimate_poses(path, writeFile)
print('Poses estimated and saved in', '\"'+writeFile+'\"', '[Time:', '{0:.2f}'.format(time.time() - start_time), 's]')
if __name__ == '__main__':
main()