-
Notifications
You must be signed in to change notification settings - Fork 1.6k
/
pose_augment.py
186 lines (139 loc) · 5.71 KB
/
pose_augment.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
import random
import math
import cv2
import numpy as np
from tensorpack.dataflow.imgaug.geometry import RotationAndCropValid
from common import CocoPart
def pose_resize_shortestedge_fixed(meta):
return pose_resize_shortestedge(meta, 368)
def pose_resize_shortestedge_random(meta):
target_size = int(400 * random.uniform(0.5, 1.1))
return pose_resize_shortestedge(meta, target_size)
def pose_resize_shortestedge(meta, target_size):
img = meta.img
# adjust image
scale = target_size * 1.0 / min(meta.height, meta.width)
if meta.height < meta.width:
newh, neww = target_size, int(scale * meta.width + 0.5)
else:
newh, neww = int(scale * meta.height + 0.5), target_size
dst = cv2.resize(img, (neww, newh), interpolation=cv2.INTER_AREA)
pw = ph = 0
if neww < 370 or newh < 370:
pw = max(0, (370 - neww) // 2)
ph = max(0, (370 - newh) // 2)
dst = cv2.copyMakeBorder(dst, ph, ph, pw, pw, cv2.BORDER_CONSTANT, value=(255, 255, 255))
# adjust meta data
adjust_joint_list = []
for joint in meta.joint_list:
adjust_joint = []
for point in joint:
if point[0] <= 0 or point[1] <= 0:
adjust_joint.append((-1, -1))
continue
adjust_joint.append((int(point[0]*scale+0.5) + pw, int(point[1]*scale+0.5) + ph))
adjust_joint_list.append(adjust_joint)
meta.joint_list = adjust_joint_list
meta.width, meta.height = neww + pw * 2, newh + ph * 2
meta.img = dst
return meta
def pose_crop_center(meta):
target_size = (368, 368)
x = (meta.width - target_size[0]) // 2 if meta.width > target_size[0] else 0
y = (meta.height - target_size[1]) // 2 if meta.height > target_size[1] else 0
return pose_crop(meta, x, y, target_size[0], target_size[1])
def pose_crop_random(meta):
target_size = (368, 368)
x = random.randrange(0, meta.width - target_size[0]) if meta.width > target_size[0] else 0
y = random.randrange(0, meta.height - target_size[1]) if meta.height > target_size[1] else 0
return pose_crop(meta, x, y, target_size[0], target_size[1])
def pose_crop(meta, x, y, w, h):
# adjust image
target_size = (w, h)
img = meta.img
resized = img[y:y+target_size[1], x:x+target_size[0], :]
# adjust meta data
adjust_joint_list = []
for joint in meta.joint_list:
adjust_joint = []
for point in joint:
if point[0] <= 0 or point[1] <= 0:
adjust_joint.append((-1, -1))
continue
new_x, new_y = point[0] - x, point[1] - y
if new_x <= 0 or new_y <= 0 or new_x > target_size[0] or new_y > target_size[1]:
adjust_joint.append((-1, -1))
continue
adjust_joint.append((new_x, new_y))
adjust_joint_list.append(adjust_joint)
meta.joint_list = adjust_joint_list
meta.width, meta.height = target_size
meta.img = resized
return meta
def pose_flip(meta):
r = random.uniform(0, 1.0)
if r > 0.5:
return meta
img = meta.img
img = cv2.flip(img, 1)
# flip meta
flip_list = [CocoPart.Nose, CocoPart.Neck, CocoPart.LShoulder, CocoPart.LElbow, CocoPart.LWrist, CocoPart.RShoulder, CocoPart.RElbow, CocoPart.RWrist,
CocoPart.LHip, CocoPart.LKnee, CocoPart.LAnkle, CocoPart.RHip, CocoPart.RKnee, CocoPart.RAnkle,
CocoPart.LEye, CocoPart.REye, CocoPart.LEar, CocoPart.REar, CocoPart.Background]
adjust_joint_list = []
for joint in meta.joint_list:
adjust_joint = []
for cocopart in flip_list:
point = joint[cocopart.value]
if point[0] <= 0 or point[1] <= 0:
adjust_joint.append((-1, -1))
continue
adjust_joint.append((meta.width - point[0], point[1]))
adjust_joint_list.append(adjust_joint)
meta.joint_list = adjust_joint_list
meta.img = img
return meta
def pose_rotation(meta):
deg = random.uniform(-40.0, 40.0)
img = meta.img
center = (img.shape[1] * 0.5, img.shape[0] * 0.5)
rot_m = cv2.getRotationMatrix2D((center[0] - 0.5, center[1] - 0.5), deg, 1)
ret = cv2.warpAffine(img, rot_m, img.shape[1::-1], flags=cv2.INTER_AREA, borderMode=cv2.BORDER_CONSTANT)
if img.ndim == 3 and ret.ndim == 2:
ret = ret[:, :, np.newaxis]
neww, newh = RotationAndCropValid.largest_rotated_rect(ret.shape[1], ret.shape[0], deg)
neww = min(neww, ret.shape[1])
newh = min(newh, ret.shape[0])
newx = int(center[0] - neww * 0.5)
newy = int(center[1] - newh * 0.5)
# print(ret.shape, deg, newx, newy, neww, newh)
img = ret[newy:newy + newh, newx:newx + neww]
# adjust meta data
adjust_joint_list = []
for joint in meta.joint_list:
adjust_joint = []
for point in joint:
if point[0] <= 0 or point[1] <= 0:
adjust_joint.append((-1, -1))
continue
x, y = _rotate_coord((meta.width, meta.height), (newx, newy), point, deg)
adjust_joint.append((x, y))
adjust_joint_list.append(adjust_joint)
meta.joint_list = adjust_joint_list
meta.width, meta.height = neww, newh
meta.img = img
return meta
def _rotate_coord(shape, newxy, point, angle):
angle = -1 * angle / 180.0 * math.pi
ox, oy = shape
px, py = point
ox /= 2
oy /= 2
qx = math.cos(angle) * (px - ox) - math.sin(angle) * (py - oy)
qy = math.sin(angle) * (px - ox) + math.cos(angle) * (py - oy)
new_x, new_y = newxy
qx += ox - new_x
qy += oy - new_y
return int(qx + 0.5), int(qy + 0.5)
def pose_to_img(meta_l):
return [meta_l[0].img, meta_l[0].get_heatmap(target_size=(92, 92)), meta_l[0].get_vectormap(target_size=(92, 92))]