-
Notifications
You must be signed in to change notification settings - Fork 1.9k
/
post_process.py
114 lines (102 loc) · 3.87 KB
/
post_process.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
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
from .image import transform_preds
from .ddd_utils import ddd2locrot
def get_pred_depth(depth):
return depth
def get_alpha(rot):
# output: (B, 8) [bin1_cls[0], bin1_cls[1], bin1_sin, bin1_cos,
# bin2_cls[0], bin2_cls[1], bin2_sin, bin2_cos]
# return rot[:, 0]
idx = rot[:, 1] > rot[:, 5]
alpha1 = np.arctan2(rot[:, 2], rot[:, 3]) + (-0.5 * np.pi)
alpha2 = np.arctan2(rot[:, 6], rot[:, 7]) + ( 0.5 * np.pi)
return alpha1 * idx + alpha2 * (1 - idx)
def ddd_post_process_2d(dets, c, s, opt):
# dets: batch x max_dets x dim
# return 1-based class det list
ret = []
include_wh = dets.shape[2] > 16
for i in range(dets.shape[0]):
top_preds = {}
dets[i, :, :2] = transform_preds(
dets[i, :, 0:2], c[i], s[i], (opt.output_w, opt.output_h))
classes = dets[i, :, -1]
for j in range(opt.num_classes):
inds = (classes == j)
top_preds[j + 1] = np.concatenate([
dets[i, inds, :3].astype(np.float32),
get_alpha(dets[i, inds, 3:11])[:, np.newaxis].astype(np.float32),
get_pred_depth(dets[i, inds, 11:12]).astype(np.float32),
dets[i, inds, 12:15].astype(np.float32)], axis=1)
if include_wh:
top_preds[j + 1] = np.concatenate([
top_preds[j + 1],
transform_preds(
dets[i, inds, 15:17], c[i], s[i], (opt.output_w, opt.output_h))
.astype(np.float32)], axis=1)
ret.append(top_preds)
return ret
def ddd_post_process_3d(dets, calibs):
# dets: batch x max_dets x dim
# return 1-based class det list
ret = []
for i in range(len(dets)):
preds = {}
for cls_ind in dets[i].keys():
preds[cls_ind] = []
for j in range(len(dets[i][cls_ind])):
center = dets[i][cls_ind][j][:2]
score = dets[i][cls_ind][j][2]
alpha = dets[i][cls_ind][j][3]
depth = dets[i][cls_ind][j][4]
dimensions = dets[i][cls_ind][j][5:8]
wh = dets[i][cls_ind][j][8:10]
locations, rotation_y = ddd2locrot(
center, alpha, dimensions, depth, calibs[0])
bbox = [center[0] - wh[0] / 2, center[1] - wh[1] / 2,
center[0] + wh[0] / 2, center[1] + wh[1] / 2]
pred = [alpha] + bbox + dimensions.tolist() + \
locations.tolist() + [rotation_y, score]
preds[cls_ind].append(pred)
preds[cls_ind] = np.array(preds[cls_ind], dtype=np.float32)
ret.append(preds)
return ret
def ddd_post_process(dets, c, s, calibs, opt):
# dets: batch x max_dets x dim
# return 1-based class det list
dets = ddd_post_process_2d(dets, c, s, opt)
dets = ddd_post_process_3d(dets, calibs)
return dets
def ctdet_post_process(dets, c, s, h, w, num_classes):
# dets: batch x max_dets x dim
# return 1-based class det dict
ret = []
for i in range(dets.shape[0]):
top_preds = {}
dets[i, :, :2] = transform_preds(
dets[i, :, 0:2], c[i], s[i], (w, h))
dets[i, :, 2:4] = transform_preds(
dets[i, :, 2:4], c[i], s[i], (w, h))
classes = dets[i, :, -1]
for j in range(num_classes):
inds = (classes == j)
top_preds[j + 1] = np.concatenate([
dets[i, inds, :4].astype(np.float32),
dets[i, inds, 4:5].astype(np.float32)], axis=1).tolist()
ret.append(top_preds)
return ret
def multi_pose_post_process(dets, c, s, h, w):
# dets: batch x max_dets x 40
# return list of 39 in image coord
ret = []
for i in range(dets.shape[0]):
bbox = transform_preds(dets[i, :, :4].reshape(-1, 2), c[i], s[i], (w, h))
pts = transform_preds(dets[i, :, 5:39].reshape(-1, 2), c[i], s[i], (w, h))
top_preds = np.concatenate(
[bbox.reshape(-1, 4), dets[i, :, 4:5],
pts.reshape(-1, 34)], axis=1).astype(np.float32).tolist()
ret.append({np.ones(1, dtype=np.int32)[0]: top_preds})
return ret