pose_utils.py
6.11 KB
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
import numpy as np
from scipy.ndimage.filters import gaussian_filter
from skimage.draw import circle, line_aa, polygon
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from collections import defaultdict
import skimage.measure, skimage.transform
import sys
LIMB_SEQ = [[1,2], [1,5], [2,3], [3,4], [5,6], [6,7], [1,8], [8,9],
[9,10], [1,11], [11,12], [12,13], [1,0], [0,14], [14,16],
[0,15], [15,17], [2,16], [5,17]]
COLORS = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
LABELS = ['nose', 'neck', 'Rsho', 'Relb', 'Rwri', 'Lsho', 'Lelb', 'Lwri',
'Rhip', 'Rkne', 'Rank', 'Lhip', 'Lkne', 'Lank', 'Leye', 'Reye', 'Lear', 'Rear']
MISSING_VALUE = -1
def map_to_cord(pose_map, threshold=0.1):
all_peaks = [[] for i in range(18)]
pose_map = pose_map[..., :18]
y, x, z = np.where(np.logical_and(pose_map == pose_map.max(axis = (0, 1)),
pose_map > threshold))
for x_i, y_i, z_i in zip(x, y, z):
all_peaks[z_i].append([x_i, y_i])
x_values = []
y_values = []
for i in range(18):
if len(all_peaks[i]) != 0:
x_values.append(all_peaks[i][0][0])
y_values.append(all_peaks[i][0][1])
else:
x_values.append(MISSING_VALUE)
y_values.append(MISSING_VALUE)
return np.concatenate([np.expand_dims(y_values, -1), np.expand_dims(x_values, -1)], axis=1)
def cords_to_map(cords, img_size, sigma=6):
result = np.zeros(img_size + cords.shape[0:1], dtype='float32')
for i, point in enumerate(cords):
if point[0] == MISSING_VALUE or point[1] == MISSING_VALUE:
continue
xx, yy = np.meshgrid(np.arange(img_size[1]), np.arange(img_size[0]))
result[..., i] = np.exp(-((yy - point[0]) ** 2 + (xx - point[1]) ** 2) / (2 * sigma ** 2))
return result
def draw_pose_from_cords(pose_joints, img_size, radius=2, draw_joints=True):
colors = np.zeros(shape=img_size + (3, ), dtype=np.uint8)
mask = np.zeros(shape=img_size, dtype=bool)
if draw_joints:
for f, t in LIMB_SEQ:
from_missing = pose_joints[f][0] == MISSING_VALUE or pose_joints[f][1] == MISSING_VALUE
to_missing = pose_joints[t][0] == MISSING_VALUE or pose_joints[t][1] == MISSING_VALUE
if from_missing or to_missing:
continue
yy, xx, val = line_aa(pose_joints[f][0], pose_joints[f][1], pose_joints[t][0], pose_joints[t][1])
colors[yy, xx] = np.expand_dims(val, 1) * 255
mask[yy, xx] = True
for i, joint in enumerate(pose_joints):
if pose_joints[i][0] == MISSING_VALUE or pose_joints[i][1] == MISSING_VALUE:
continue
yy, xx = circle(joint[0], joint[1], radius=radius, shape=img_size)
colors[yy, xx] = COLORS[i]
mask[yy, xx] = True
return colors, mask
def draw_pose_from_map(pose_map, threshold=0.1, **kwargs):
cords = map_to_cord(pose_map, threshold=threshold)
return draw_pose_from_cords(cords, pose_map.shape[:2], **kwargs)
def load_pose_cords_from_strings(y_str, x_str):
y_cords = json.loads(y_str)
x_cords = json.loads(x_str)
return np.concatenate([np.expand_dims(y_cords, -1), np.expand_dims(x_cords, -1)], axis=1)
def mean_inputation(X):
X = X.copy()
for i in range(X.shape[1]):
for j in range(X.shape[2]):
val = np.mean(X[:, i, j][X[:, i, j] != -1])
X[:, i, j][X[:, i, j] == -1] = val
return X
def draw_legend():
handles = [mpatches.Patch(color=np.array(color) / 255.0, label=name) for color, name in zip(COLORS, LABELS)]
plt.legend(handles=handles, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
def produce_ma_mask(kp_array, img_size, point_radius=4):
from skimage.morphology import dilation, erosion, square
mask = np.zeros(shape=img_size, dtype=bool)
limbs = [[2,3], [2,6], [3,4], [4,5], [6,7], [7,8], [2,9], [9,10],
[10,11], [2,12], [12,13], [13,14], [2,1], [1,15], [15,17],
[1,16], [16,18], [2,17], [2,18], [9,12], [12,6], [9,3], [17,18]]
limbs = np.array(limbs) - 1
for f, t in limbs:
from_missing = kp_array[f][0] == MISSING_VALUE or kp_array[f][1] == MISSING_VALUE
to_missing = kp_array[t][0] == MISSING_VALUE or kp_array[t][1] == MISSING_VALUE
if from_missing or to_missing:
continue
norm_vec = kp_array[f] - kp_array[t]
norm_vec = np.array([-norm_vec[1], norm_vec[0]])
norm_vec = point_radius * norm_vec / np.linalg.norm(norm_vec)
vetexes = np.array([
kp_array[f] + norm_vec,
kp_array[f] - norm_vec,
kp_array[t] - norm_vec,
kp_array[t] + norm_vec
])
yy, xx = polygon(vetexes[:, 0], vetexes[:, 1], shape=img_size)
mask[yy, xx] = True
for i, joint in enumerate(kp_array):
if kp_array[i][0] == MISSING_VALUE or kp_array[i][1] == MISSING_VALUE:
continue
yy, xx = circle(joint[0], joint[1], radius=point_radius, shape=img_size)
mask[yy, xx] = True
mask = dilation(mask, square(5))
mask = erosion(mask, square(5))
return mask
if __name__ == "__main__":
import pandas as pd
from skimage.io import imread
import pylab as plt
import os
i = 5
df = pd.read_csv('data/market-annotation-train.csv', sep=':')
for index, row in df.iterrows():
pose_cords = load_pose_cords_from_strings(row['keypoints_y'], row['keypoints_x'])
colors, mask = draw_pose_from_cords(pose_cords, (128, 64))
mmm = produce_ma_mask(pose_cords, (128, 64)).astype(float)[..., np.newaxis].repeat(3, axis=-1)
# print mmm.shape
print(mmm.shape)
img = imread('data/market-dataset/train/' + row['name'])
mmm[mask] = colors[mask]
# print (mmm)
print(mmm)
plt.subplot(1, 1, 1)
plt.imshow(mmm)
plt.show()