-
Notifications
You must be signed in to change notification settings - Fork 0
/
data_utils.py
126 lines (101 loc) · 3.1 KB
/
data_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
import pickle
import os
import numpy as np
def read_pickle(pkl_path):
with open(pkl_path, 'rb') as f:
return pickle.load(f)
def save_pickle(data, pkl_path):
os.system('mkdir -p {}'.format(os.path.dirname(pkl_path)))
with open(pkl_path, 'wb') as f:
pickle.dump(data, f)
def project(xyz, K, RT):
"""
xyz: [N, 3]
K: [3, 3]
RT: [3, 4]
"""
xyz = np.dot(xyz, RT[:, :3].T) + RT[:, 3:].T
xyz = np.dot(xyz, K.T)
xy = xyz[:, :2] / xyz[:, 2:]
return xy
def get_pose_mat(pose):
position, orientation = pose
if len(orientation) == 4:
from transforms3d.quaternions import quat2mat
x, y, z, w = orientation
quat_new = [w, x, y, z]
mat_R = quat2mat(quat_new)
elif len(orientation) == 3:
from transforms3d.euler import euler2mat
x, y, z = orientation
mat_R = np.dot(np.dot(np.array(euler2mat(x, 0, 0)), np.array(euler2mat(0, y, 0))), euler2mat(0, 0, z))
else:
raise NotImplementedError
pose = np.eye(4, 4)
pose[:3, :3] = mat_R
pose[:3, 3] = position
return pose
def draw_box2d(box2d, image=None, plt=None):
x_min, y_min, x_max, y_max = box2d
points = [
[x_min, y_min],
[x_min, y_max],
[x_max, y_max],
[x_max, y_min]
]
lines = [
[0, 1],
[1, 2],
[2, 3],
[3, 0]
]
for id, line in enumerate(lines):
pt1 = points[line[0]]
pt2 = points[line[1]]
if image is not None:
import cv2
pt1 = (int(pt1[0]), int(pt1[1]))
pt2 = (int(pt2[0]), int(pt2[1]))
cv2.line(image, pt1, pt2, [0, 0, 255], 3)
if plt is not None:
plt.plot([pt1[0], pt2[0]], [pt1[1], pt2[1]])
def draw_box3d(box3d_proj, image=None, plt=None):
lines = [
[1, 2],
[1, 3],
[1, 5],
[2, 6],
[2, 4],
[4, 3],
[4, 8],
[8, 6],
[8, 7],
[5, 7],
[3, 7],
[5, 6]
]
# orange, blue, green, purple red, pink, caffe
color_list = [(47, 92, 247), (68, 27, 203), (221, 168, 81), (62, 129, 27), (129, 51, 111), (27, 177, 255),
(56, 67, 115), (47, 92, 247), (68, 27, 203), (221, 168, 81), (62, 129, 27), (27, 177, 255)]
for id, line in enumerate(lines):
pt1 = box3d_proj[line[0] - 1, :]
pt2 = box3d_proj[line[1] - 1, :]
if image is not None:
import cv2
pt1 = (int(pt1[0]), int(pt1[1]))
pt2 = (int(pt2[0]), int(pt2[1]))
cv2.line(image, pt1, pt2, color_list[id], 3)
if plt is not None:
plt.plot([pt1[0], pt2[0]], [pt1[1], pt2[1]])
def get_bbox2d(corner2d, w, h):
x_max = int(min(np.max(corner2d[:, 0]), w))
x_min = int(max(np.min(corner2d[:, 0]), 0))
y_max = int(min(np.max(corner2d[:, 1]), h))
y_min = int(max(np.min(corner2d[:, 1]), 0))
return [x_min, y_min, x_max, y_max]
def check_kp_bound(fps2d, w, h):
count = 0
for point in fps2d:
if point[0] < 0 or point[0] > w or point[1] < 0 or point[1] > h:
count += 1
return count