-
Notifications
You must be signed in to change notification settings - Fork 11
/
Copy pathkeypoint2img_withhand.py
204 lines (182 loc) · 8.69 KB
/
keypoint2img_withhand.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
import os.path
from PIL import Image
import numpy as np
import json
import glob
from scipy.optimize import curve_fit
import warnings
import cv2
def func(x, a, b, c):
return a * x**2 + b * x + c
def linear(x, a, b):
return a * x + b
def setColor(im, yy, xx, color):
if len(im.shape) == 3:
if (im[yy, xx] == 0).all():
im[yy, xx, 0], im[yy, xx, 1], im[yy, xx, 2] = color[0], color[1], color[2]
else:
im[yy, xx, 0] = ((im[yy, xx, 0].astype(float) + color[0]) / 2).astype(np.uint8)
im[yy, xx, 1] = ((im[yy, xx, 1].astype(float) + color[1]) / 2).astype(np.uint8)
im[yy, xx, 2] = ((im[yy, xx, 2].astype(float) + color[2]) / 2).astype(np.uint8)
else:
im[yy, xx] = color[0]
def drawEdge(im, x, y, bw=1, color=(255,255,255), draw_end_points=False):
if x is not None and x.size:
h, w = im.shape[0], im.shape[1]
# edge
for i in range(-bw, bw):
for j in range(-bw, bw):
yy = np.maximum(0, np.minimum(h-1, y+i))
xx = np.maximum(0, np.minimum(w-1, x+j))
setColor(im, yy, xx, color)
# edge endpoints
if draw_end_points:
for i in range(-bw*2, bw*2):
for j in range(-bw*2, bw*2):
if (i**2) + (j**2) < (4 * bw**2):
yy = np.maximum(0, np.minimum(h-1, np.array([y[0], y[-1]])+i))
xx = np.maximum(0, np.minimum(w-1, np.array([x[0], x[-1]])+j))
setColor(im, yy, xx, color)
def interpPoints(x, y):
if abs(x[:-1] - x[1:]).max() < abs(y[:-1] - y[1:]).max():
curve_y, curve_x = interpPoints(y, x)
if curve_y is None:
return None, None
else:
with warnings.catch_warnings():
warnings.simplefilter("ignore")
if len(x) < 3:
popt, _ = curve_fit(linear, x, y)
else:
popt, _ = curve_fit(func, x, y)
if abs(popt[0]) > 1:
return None, None
if x[0] > x[-1]:
x = list(reversed(x))
y = list(reversed(y))
curve_x = np.linspace(x[0], x[-1], (x[-1]-x[0]))
if len(x) < 3:
curve_y = linear(curve_x, *popt)
else:
curve_y = func(curve_x, *popt)
return curve_x.astype(int), curve_y.astype(int)
def read_keypoints(json_input, size, random_drop_prob=0, remove_face_labels=False, basic_point_only=False):
with open(json_input, encoding='utf-8') as f:
keypoint_dicts = json.loads(f.read())["people"]
edge_lists = define_edge_lists(basic_point_only)
w, h = size
pose_img = np.zeros((h, w, 3), np.uint8)
for keypoint_dict in keypoint_dicts:
pose_pts = np.array(keypoint_dict["pose_keypoints_2d"]).reshape(25, 3)
face_pts = np.array(keypoint_dict["face_keypoints_2d"]).reshape(70, 3)
hand_pts_l = np.array(keypoint_dict["hand_left_keypoints_2d"]).reshape(21, 3)
hand_pts_r = np.array(keypoint_dict["hand_right_keypoints_2d"]).reshape(21, 3)
pts = [extract_valid_keypoints(pts, edge_lists) for pts in [pose_pts, face_pts, hand_pts_l, hand_pts_r]]
pose_img += connect_keypoints(pts, edge_lists, size, random_drop_prob, remove_face_labels, basic_point_only)
return pose_img
def extract_valid_keypoints(pts, edge_lists):
pose_edge_list, _, hand_edge_list, _, face_list = edge_lists
p = pts.shape[0]
thre = 0.1 if p == 70 else 0.01
output = np.zeros((p, 2))
if p == 70: # face
for edge_list in face_list:
for edge in edge_list:
if (pts[edge, 2] > thre).all():
output[edge, :] = pts[edge, :2]
elif p == 21: # hand
for edge in hand_edge_list:
if (pts[edge, 2] > thre).all():
output[edge, :] = pts[edge, :2]
else: # pose
valid = (pts[:, 2] > thre)
output[valid, :] = pts[valid, :2]
return output
def connect_keypoints(pts, edge_lists, size, random_drop_prob, remove_face_labels, basic_point_only):
pose_pts, face_pts, hand_pts_l, hand_pts_r = pts
w, h = size
output_edges = np.zeros((h, w, 3), np.uint8)
pose_edge_list, pose_color_list, hand_edge_list, hand_color_list, face_list = edge_lists
if random_drop_prob > 0 and remove_face_labels:
# add random noise to keypoints
pose_pts[[0,15,16,17,18], :] += 5 * np.random.randn(5,2)
face_pts[:,0] += 2 * np.random.randn()
face_pts[:,1] += 2 * np.random.randn()
### pose
for i, edge in enumerate(pose_edge_list):
x, y = pose_pts[edge, 0], pose_pts[edge, 1]
if (np.random.rand() > random_drop_prob) and (0 not in x):
curve_x, curve_y = interpPoints(x, y)
drawEdge(output_edges, curve_x, curve_y, bw=3, color=pose_color_list[i], draw_end_points=True)
if not basic_point_only:
### hand
for hand_pts in [hand_pts_l, hand_pts_r]: # for left and right hand
if np.random.rand() > random_drop_prob:
for i, edge in enumerate(hand_edge_list): # for each finger
for j in range(0, len(edge)-1): # for each part of the finger
sub_edge = edge[j:j+2]
x, y = hand_pts[sub_edge, 0], hand_pts[sub_edge, 1]
if 0 not in x:
line_x, line_y = interpPoints(x, y)
drawEdge(output_edges, line_x, line_y, bw=3, color=hand_color_list[i], draw_end_points=True)
### face
edge_len = 2
if (np.random.rand() > random_drop_prob):
for edge_list in face_list:
for edge in edge_list:
for i in range(0, max(1, len(edge)-1), edge_len-1):
sub_edge = edge[i:i+edge_len]
x, y = face_pts[sub_edge, 0], face_pts[sub_edge, 1]
if 0 not in x:
curve_x, curve_y = interpPoints(x, y)
drawEdge(output_edges, curve_x, curve_y, bw=2, draw_end_points=False)
###wrist points
#cv2.circle(output_edges, (int(pose_pts[7, 0]), int(pose_pts[7, 1])), 10, (0, 255, 0), -1)
#cv2.circle(output_edges, (int(pose_pts[4, 0]), int(pose_pts[4, 1])), 10, (255, 0, 0), -1)
cv2.circle(output_edges, (int(hand_pts_l[9, 0]), int(hand_pts_l[9, 1])), 8, (0, 255, 0), -1)
cv2.circle(output_edges, (int(hand_pts_r[9, 0]), int(hand_pts_r[9, 1])), 8, (255, 0, 0), -1)
return output_edges
def define_edge_lists(basic_point_only):
### pose
pose_edge_list = []
pose_color_list = []
#if not basic_point_only:
#pose_edge_list += [[17, 15], [15, 0], [ 0, 16], [16, 18]] # head
#pose_color_list += [[153, 0,153], [153, 0,102], [102, 0,153], [ 51, 0,153]]
pose_edge_list += [
[ 0, 1], [ 1, 8], # body
[ 1, 2], [ 2, 3], [ 3, 4], # right arm
[ 1, 5], [ 5, 6], [ 6, 7], # left arm
[ 8, 9], #[ 9, 10], [10, 11], [11, 24], [11, 22], [22, 23], # right leg
[ 8, 12]#, [12, 13], [13, 14], [14, 21], [14, 19], [19, 20] # left leg
]
pose_color_list += [
[153, 0, 51], [153, 0, 0],
[153, 51, 0], [153,102, 0], [153,153, 0],
[102,153, 0], [ 51,153, 0], [ 0,153, 0],
[ 0,153, 51], [ 0,153,102], [ 0,153,153], [ 0,153,153], [ 0,153,153], [ 0,153,153],
[ 0,102,153], [ 0, 51,153], [ 0, 0,153], [ 0, 0,153], [ 0, 0,153], [ 0, 0,153]
]
### hand
hand_edge_list = [
[0, 1, 2, 3, 4],
[0, 5, 6, 7, 8],
[0, 9, 10, 11, 12],
[0, 13, 14, 15, 16],
[0, 17, 18, 19, 20]
]
hand_color_list = [
[204,0,0], [163,204,0], [0,204,82], [0,82,204], [163,0,204]
]
### face
face_list = [
[range(0, 17)], # face
[range(17, 22)], # left eyebrow
[range(22, 27)], # right eyebrow
[range(27, 31), range(31, 36)], # nose
[[36,37,38,39], [39,40,41,36]], # left eye
[[42,43,44,45], [45,46,47,42]], # right eye
[range(48, 55), [54,55,56,57,58,59,48]], # mouth outer
[range(60, 65), [64, 65, 66, 67, 60]], # mouth inner
]
return pose_edge_list, pose_color_list, hand_edge_list, hand_color_list, face_list