-
Notifications
You must be signed in to change notification settings - Fork 128
/
Copy pathvis_label.py
95 lines (78 loc) · 3.77 KB
/
vis_label.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
import os
import tqdm
import argparse
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from data.dataset import HDMapNetDataset, CAMS
from data.utils import get_proj_mat, perspective
from data.image import denormalize_img
def vis_label(dataroot, version, xbound, ybound):
data_conf = {
'image_size': (900, 1600),
'xbound': xbound,
'ybound': ybound
}
color_map = np.random.randint(0, 256, (256, 3))
color_map[0] = np.array([0, 0, 0])
colors_plt = ['r', 'b', 'g']
dataset = HDMapNetDataset(version=version, dataroot=dataroot, data_conf=data_conf, is_train=False)
gt_path = os.path.join(dataroot, 'samples', 'GT')
if not os.path.exists(gt_path):
os.mkdir(gt_path)
car_img = Image.open('icon/car.png')
for idx in tqdm.tqdm(range(dataset.__len__())):
rec = dataset.nusc.sample[idx]
imgs, trans, rots, intrins, post_trans, post_rots = dataset.get_imgs(rec)
vectors = dataset.get_vectors(rec)
lidar_top_path = dataset.nusc.get_sample_data_path(rec['data']['LIDAR_TOP'])
base_path = lidar_top_path.split('/')[-1].replace('__LIDAR_TOP__', '_').split('.')[0]
base_path = os.path.join(gt_path, base_path)
if not os.path.exists(base_path):
os.mkdir(base_path)
plt.figure(figsize=(4, 2))
plt.xlim(-30, 30)
plt.ylim(-15, 15)
plt.axis('off')
for vector in vectors:
pts, pts_num, line_type = vector['pts'], vector['pts_num'], vector['type']
pts = pts[:pts_num]
x = np.array([pt[0] for pt in pts])
y = np.array([pt[1] for pt in pts])
plt.quiver(x[:-1], y[:-1], x[1:] - x[:-1], y[1:] - y[:-1], scale_units='xy', angles='xy', scale=1, color=colors_plt[line_type])
plt.imshow(car_img, extent=[-1.5, 1.5, -1.2, 1.2])
map_path = os.path.join(base_path, 'MAP.png')
plt.savefig(map_path, bbox_inches='tight', dpi=400)
plt.close()
for img, intrin, rot, tran, cam in zip(imgs, intrins, rots, trans, CAMS):
img = denormalize_img(img)
P = get_proj_mat(intrin, rot, tran)
plt.figure(figsize=(9, 16))
fig = plt.imshow(img)
fig.axes.get_xaxis().set_visible(False)
fig.axes.get_yaxis().set_visible(False)
plt.xlim(1600, 0)
plt.ylim(900, 0)
plt.axis('off')
for vector in vectors:
pts, pts_num, line_type = vector['pts'], vector['pts_num'], vector['type']
pts = pts[:pts_num]
zeros = np.zeros((pts_num, 1))
ones = np.ones((pts_num, 1))
world_coords = np.concatenate([pts, zeros, ones], axis=1).transpose(1, 0)
pix_coords = perspective(world_coords, P)
x = np.array([pts[0] for pts in pix_coords])
y = np.array([pts[1] for pts in pix_coords])
plt.quiver(x[:-1], y[:-1], x[1:] - x[:-1], y[1:] - y[:-1], scale_units='xy',
angles='xy', scale=1, color=colors_plt[line_type])
cam_path = os.path.join(base_path, f'{cam}.png')
plt.savefig(cam_path, bbox_inches='tight', pad_inches=0, dpi=400)
plt.close()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Local HD Map Demo.')
parser.add_argument('dataroot', type=str, default='dataset/nuScenes/')
parser.add_argument('--version', type=str, default='v1.0-mini', choices=['v1.0-trainval', 'v1.0-mini'])
parser.add_argument("--xbound", nargs=3, type=float, default=[-30.0, 30.0, 0.15])
parser.add_argument("--ybound", nargs=3, type=float, default=[-15.0, 15.0, 0.15])
args = parser.parse_args()
vis_label(args.dataroot, args.version, args.xbound, args.ybound)