-
Notifications
You must be signed in to change notification settings - Fork 1
/
utils.py
187 lines (153 loc) · 5.58 KB
/
utils.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
import math
import paddle
import shutil
import time
import os
import random
from easydict import EasyDict as edict
import yaml
import numpy as np
class AverageMeter(object):
""" Computes ans stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0.
self.avg = 0.
self.sum = 0.
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def to_torch(ndarray):
if type(ndarray).__module__ == 'numpy':
return paddle.to_tensor(ndarray)
elif not paddle.is_tensor(ndarray):
raise ValueError("Cannot convert {} to torch tensor"
.format(type(ndarray)))
return ndarray
def get_transform(center, scale, res, rot=0):
"""
General image processing functions
"""
# Generate transformation matrix
h = 200 * scale
t = np.zeros((3, 3))
t[0, 0] = float(res[1]) / h
t[1, 1] = float(res[0]) / h
t[0, 2] = res[1] * (-float(center[0]) / h + .5)
t[1, 2] = res[0] * (-float(center[1]) / h + .5)
t[2, 2] = 1
if not rot == 0:
rot = -rot # To match direction of rotation from cropping
rot_mat = np.zeros((3,3))
rot_rad = rot * np.pi / 180
sn,cs = np.sin(rot_rad), np.cos(rot_rad)
rot_mat[0,:2] = [cs, -sn]
rot_mat[1,:2] = [sn, cs]
rot_mat[2,2] = 1
# Need to rotate around center
t_mat = np.eye(3)
t_mat[0,2] = -res[1]/2
t_mat[1,2] = -res[0]/2
t_inv = t_mat.copy()
t_inv[:2,2] *= -1
t = np.dot(t_inv,np.dot(rot_mat,np.dot(t_mat,t)))
return t
def transform(pt, center, scale, res, invert=0, rot=0):
# Transform pixel location to different reference
t = get_transform(center, scale, res, rot=rot)
if invert:
t = np.linalg.inv(t)
new_pt = np.array([pt[0] - 1, pt[1] - 1, 1.]).T
new_pt = np.dot(t, new_pt)
return new_pt[:2].astype(int) + 1
def transform_preds(coords, center, scale, res):
# size = coords.size()
# coords = coords.view(-1, coords.size(-1))
# print(coords.size())
for p in range(coords.size(0)):
coords[p, 0:2] = to_torch(transform(coords[p, 0:2], center, scale, res, 1, 0))
return coords
def get_preds(scores):
''' get predictions from score maps in torch Tensor
return type: torch.LongTensor
'''
assert scores.dim() == 4, 'Score maps should be 4-dim'
maxval, idx = paddle.max(scores.reshape(scores.shape[0], scores.shape[1], -1), 2)
maxval = maxval.reshape(scores.shape[0], scores.shape[1], 1)
idx = idx.reshape(scores.shape[0], scores.shape[1], 1) + 1
preds = idx.repeat(1, 1, 2).float()
preds[:,:,0] = (preds[:,:,0] - 1) % scores.size(3) + 1
preds[:,:,1] = paddle.floor((preds[:,:,1] - 1) / scores.shape[3]) + 1
pred_mask = maxval.gt(0).repeat(1, 1, 2).float()
preds *= pred_mask
return preds
def calc_dists(preds, target, normalize):
preds = preds.float()
target = target.float()
dists = paddle.zeros((preds.shape[1], preds.shape[0]))
for n in range(preds.shape[0]):
for c in range(preds.shape[1]):
if target[n,c,0] > 1 and target[n, c, 1] > 1:
dists[c, n] = paddle.dist(preds[n,c,:], target[n,c,:])/normalize[n]
else:
dists[c, n] = -1
return dists
def dist_acc(dist, thr=0.5):
''' Return percentage below threshold while ignoring values with a -1 '''
dist = dist[dist != -1]
if len(dist) > 0:
return 1.0 * (dist < thr).sum().item() / len(dist)
else:
return -1
def accuracy(output, target, idxs, thr=0.5):
''' Calculate accuracy according to PCK, but uses ground truth heatmap rather than x,y locations
First value to be returned is average accuracy across 'idxs', followed by individual accuracies
'''
preds = get_preds(output)
gts = get_preds(target)
norm = paddle.ones(preds.shape[0])*output.shape[3]/10
dists = calc_dists(preds, gts, norm)
acc = paddle.zeros([len(idxs)+1])
avg_acc = 0
cnt = 0
for i in range(len(idxs)):
acc[i+1] = dist_acc(dists[idxs[i]-1])
if acc[i+1] >= 0:
avg_acc = avg_acc + acc[i+1]
cnt += 1
if cnt != 0:
acc[0] = avg_acc / cnt
return acc
def final_preds(output, center, scale, res):
coords = get_preds(output) # float type
# pose-processing
for n in range(coords.shape[0]):
for p in range(coords.shape[1]):
hm = output[n][p]
px = int(math.floor(coords[n][p][0]))
py = int(math.floor(coords[n][p][1]))
if px > 1 and px < res[0] and py > 1 and py < res[1]:
diff = paddle.Tensor([hm[py - 1][px] - hm[py - 1][px - 2], hm[py][px - 1]-hm[py - 2][px - 1]])
coords[n][p] += diff.sign() * .25
coords += 0.5
preds = coords.clone()
# Transform back
for i in range(coords.shape[0]):
preds[i] = transform_preds(coords[i], center[i], scale[i], res)
if preds.dim() < 3:
preds = preds.view(1, preds.size())
return preds
def save_checkpoint(state, is_best, filename='checkpoint.pdparams.tar'):
paddle.save(state, './output/' + filename + '_latest.pdparams.tar')
if is_best:
shutil.copyfile('./output/' + filename + '_latest.pdparams.tar', './output/' + filename + '_best.pdparams.tar')
def Config(filename):
with open(filename, 'r') as f:
parser = edict(yaml.load(f))
for x in parser:
print('{}: {}'.format(x, parser[x]))
return parser