├── assets
├── multiagent_21.gif
└── multitraversal_61.gif
├── readme.md
└── visualize
├── lidar_cam_viz_agent.py
├── lidar_cam_viz_traversal.py
└── lidar_viz.py
/assets/multiagent_21.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ai4ce/MARS/dd1baed493cef7aa99544417bf7632bcbfa59a26/assets/multiagent_21.gif
--------------------------------------------------------------------------------
/assets/multitraversal_61.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ai4ce/MARS/dd1baed493cef7aa99544417bf7632bcbfa59a26/assets/multitraversal_61.gif
--------------------------------------------------------------------------------
/readme.md:
--------------------------------------------------------------------------------
1 | # Multiagent Multitraversal Multimodal Self-Driving: Open MARS Dataset
2 | [Yiming Li](https://roboticsyimingli.github.io/),
3 | [Zhiheng Li](https://zl3466.github.io/),
4 | [Nuo Chen](),
5 | [Moonjun Gong](https://moonjungong.github.io/),
6 | [Zonglin Lyu](),
7 | [Zehong Wang](),
8 | [Peili Jiang](),
9 | [Chen Feng](https://engineering.nyu.edu/faculty/chen-feng)
10 |
11 | [Paper](https://arxiv.org/abs/2406.09383)
12 |
13 | [Tutorial](#tutorial)
14 |
15 | Checkout our [project website](https://ai4ce.github.io/MARS/) for more demo videos.
16 | Codes to reproduce the videos are available in `/visualize` folder of `main` branch.
17 |
18 | 
19 |
20 | # Multiagent
21 | 
22 |
23 | # Multitraversal
24 | 
25 |
26 |
27 |
28 |
29 |
30 | # News
31 |
32 | - [2024/06] Both Multiagent and Multitraversal subsets are now available for download on [huggingface](https://huggingface.co/datasets/ai4ce/MARS).
33 |
34 | - [2024/06]The preprint version is available on [arXiv]([https://huggingface.co/datasets/ai4ce/MARS](https://arxiv.org/abs/2406.09383)).
35 |
36 | - [2024/02] Our paper has been accepted on CVPR 2024 🎉🎉🎉
37 |
38 |
39 |
40 |
41 |
42 | # Abstract
43 | In collaboration with the self-driving company [May Mobility](https://maymobility.com/), we present the MARS dataset which unifies scenarios that enable multiagent, multitraversal, and multimodal autonomous vehicle research.
44 |
45 | MARS is collected with a fleet of autonomous vehicles driving within a certain geographical area. Each vehicle has its own route and different vehicles may appear at nearby locations. Each vehicle is equipped with a LiDAR and surround-view RGB cameras.
46 |
47 | We curate two subsets in MARS: one facilitates collaborative driving with multiple vehicles simultaneously present at the same location, and the other enables memory retrospection through asynchronous traversals of the same location by multiple vehicles. We conduct experiments in place recognition and neural reconstruction. More importantly, MARS introduces new research opportunities and challenges such as multitraversal 3D reconstruction, multiagent perception, and unsupervised object discovery.
48 |
49 |
50 | ## Our dataset uses the same structure as the [NuScenes](https://www.nuscenes.org/nuscenes) Dataset:
51 |
52 | - Multitraversal: each location is saved as one NuScenes object, and each traversal is one scene.
53 | - Multiagent: the whole set is a NuScenes object, and each multiagent encounter is one scene.
54 |
55 |
56 |
57 | # License
58 | [CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/)
59 |
60 |
61 |
62 | # Bibtex
63 |
64 | ```
65 | @InProceedings{Li_2024_CVPR,
66 | author = {Li, Yiming and Li, Zhiheng and Chen, Nuo and Gong, Moonjun and Lyu, Zonglin and Wang, Zehong and Jiang, Peili and Feng, Chen},
67 | title = {Multiagent Multitraversal Multimodal Self-Driving: Open MARS Dataset},
68 | booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
69 | month = {June},
70 | year = {2024},
71 | pages = {22041-22051}
72 | }
73 | ```
74 |
75 |
76 |
77 | # Tutorial
78 | This tutorial explains how the NuScenes structure works in our dataset, including how you may access a scene and query its samples of sensor data.
79 |
80 | - [Devkit Initialization](#initialization)
81 | - [Multitraversal](#load-multitraversal)
82 | - [Multiagent](#load-multiagent)
83 | - [Scene](#scene)
84 | - [Sample](#sample)
85 | - [Sample Data](#sample-data)
86 | - [Sensor Names](#sensor-names)
87 | - [Camera](#camera-data)
88 | - [LiDAR](#lidar-data)
89 | - [IMU](#imu-data)
90 | - [Ego & Sensor Pose](#vehicle-and-sensor-pose)
91 | - [LiDAR-Image projection](#lidar-image-projection)
92 |
93 |
94 |
95 | ## Initialization
96 | First, install `nuscenes-devkit` following NuScenes's repo tutorial, [Devkit setup section](https://github.com/nutonomy/nuscenes-devkit?tab=readme-ov-file#devkit-setup). The easiest way is install via pip:
97 | ```
98 | pip install nuscenes-devkit
99 | ```
100 |
101 | Import NuScenes devkit:
102 | ```
103 | from nuscenes.nuscenes import NuScenes
104 | ```
105 |
106 | #### Load Multitraversal
107 | loading data of location 10:
108 | ```
109 | # The "version" variable is the name of the folder holding all .json metadata tables.
110 | location = 10
111 | nusc = NuScenes(version='v1.0', dataroot=f'/MARS_multitraversal/{location}', verbose=True)
112 | ```
113 |
114 | #### Load Multiagent
115 | loading data for the full set:
116 | ```
117 | nusc = NuScenes(version='v1.0', dataroot=f'/MARS_multiagent', verbose=True)
118 | ```
119 |
120 |
121 |
122 | ## Scene
123 | To see all scenes in one set (one location of the Multitraversal set, or the whole Multiagent set):
124 | ```
125 | print(nusc.scene)
126 | ```
127 | Output:
128 | ```
129 | [{'token': '97hitl8ya1335v8zkixvsj3q69tgx801', 'nbr_samples': 611, 'first_sample_token': 'udrq868482482o88p9r2n8b86li7cfxx', 'last_sample_token': '7s5ogk8m9id7apixkqoh3rep0s9113xu', 'name': '2023_10_04_scene_3_maisy', 'intersection': 10, 'err_max': 20068.00981996727},
130 | {'token': 'o858jv3a464383gk9mm8at71ai994d3n', 'nbr_samples': 542, 'first_sample_token': '933ho5988jo3hu848b54749x10gd7u14', 'last_sample_token': 'os54se39x1px2ve12x3r1b87e0d7l1gn', 'name': '2023_10_04_scene_4_maisy', 'intersection': 10, 'err_max': 23959.357933579337},
131 | {'token': 'xv2jkx6m0o3t044bazyz9nwbe5d5i7yy', 'nbr_samples': 702, 'first_sample_token': '8rqb40c919d6n5cd553c3j01v178k28m', 'last_sample_token': 'skr79z433oyi6jljr4nx7ft8c42549nn', 'name': '2023_10_04_scene_6_mike', 'intersection': 10, 'err_max': 27593.048433048432},
132 | {'token': '48e90c7dx401j97391g6549zmljbg0hk', 'nbr_samples': 702, 'first_sample_token': 'ui8631xb2in5la133319c5301wvx1fib', 'last_sample_token': 'xrns1rpma4p00hf39305ckol3p91x59w', 'name': '2023_10_04_scene_9_mike', 'intersection': 10, 'err_max': 24777.237891737892},
133 | ...
134 | ]
135 |
136 | ```
137 |
138 | The scenes can then be retrieved by indexing:
139 | ```
140 | num_of_scenes = len(nusc.scene)
141 | my_scene = nusc.scene[0] # scene at index 0, which is the first scene of this location
142 | print(first_scene)
143 | ```
144 | Output:
145 | ```
146 | {'token': '97hitl8ya1335v8zkixvsj3q69tgx801',
147 | 'nbr_samples': 611,
148 | 'first_sample_token': 'udrq868482482o88p9r2n8b86li7cfxx',
149 | 'last_sample_token': '7s5ogk8m9id7apixkqoh3rep0s9113xu',
150 | 'name': '2023_10_04_scene_3_maisy',
151 | 'intersection': 10,
152 | 'err_max': 20068.00981996727}
153 | ```
154 | - `nbr_samples`: number of samples (frames) of this scene.
155 | - `name`: name of the scene, including its date and name of the vehicle it is from (in this example, the data is from Oct. 4th 2023, vehicle maisy).
156 | - `intersection`: location index.
157 | - `err_max`: maximum time difference (in millisecond) between camera images of a same frame in this scene.
158 |
159 |
160 |
161 | ## Sample
162 | Get the first sample (frame) of one scene:
163 | ```
164 | first_sample_token = my_scene['first_sample_token'] # get sample token
165 | my_sample = nusc.get('sample', first_sample_token) # get sample metadata
166 | print(my_sample)
167 | ```
168 |
169 | Output:
170 | ```
171 | {'token': 'udrq868482482o88p9r2n8b86li7cfxx',
172 | 'timestamp': 1696454482883182,
173 | 'prev': '',
174 | 'next': 'v15b2l4iaq1x0abxr45jn6bi08j72i01',
175 | 'scene_token': '97hitl8ya1335v8zkixvsj3q69tgx801',
176 | 'data': {
177 | 'CAM_FRONT_CENTER': 'q9e0pgk3wiot983g4ha8178zrnr37m50',
178 | 'CAM_FRONT_LEFT': 'c13nf903o913k30rrz33b0jq4f0z7y2d',
179 | 'CAM_FRONT_RIGHT': '67ydh75sam2dtk67r8m3bk07ba0lz3ib',
180 | 'CAM_BACK_CENTER': '1n09qfm9vw65xpohjqgji2g58459gfuq',
181 | 'CAM_SIDE_LEFT': '14up588181925s8bqe3pe44d60316ey0',
182 | 'CAM_SIDE_RIGHT': 'x95k7rvhmxkndcj8mc2821c1cs8d46y5',
183 | 'LIDAR_FRONT_CENTER': '13y90okaf208cqqy1v54z87cpv88k2qy',
184 | 'IMU_TOP': 'to711a9v6yltyvxn5653cth9w2o493z4'
185 | },
186 | 'anns': []}
187 | ```
188 | - `prev`: token of the previous sample.
189 | - `next`': token of the next sample.
190 | - `data`: dict of data tokens of this sample's sensor data.
191 | - `anns`: empty as we do not have annotation data at this moment.
192 |
193 |
194 |
195 | ## Sample Data
196 | ### Sensor Names
197 | Our sensor names are different from NuScenes' sensor names. It is important that you use the correct name when querying sensor data. Our sensor names are:
198 | ```
199 | ['CAM_FRONT_CENTER',
200 | 'CAM_FRONT_LEFT',
201 | 'CAM_FRONT_RIGHT',
202 | 'CAM_BACK_CENTER',
203 | 'CAM_SIDE_LEFT',
204 | 'CAM_SIDE_RIGHT',
205 | 'LIDAR_FRONT_CENTER',
206 | 'IMU_TOP']
207 | ```
208 |
209 | ---
210 | ### Camera Data
211 | All image data are already undistorted.
212 |
213 | To load a piece data, we start with querying its `sample_data` dictionary object from the metadata:
214 | ```
215 | sensor = 'CAM_FRONT_CENTER'
216 | sample_data_token = my_sample['data'][sensor]
217 | FC_data = nusc.get('sample_data', sample_data_token)
218 | print(FC_data)
219 | ```
220 | Output:
221 | ```
222 | {'token': 'q9e0pgk3wiot983g4ha8178zrnr37m50',
223 | 'sample_token': 'udrq868482482o88p9r2n8b86li7cfxx',
224 | 'ego_pose_token': 'q9e0pgk3wiot983g4ha8178zrnr37m50',
225 | 'calibrated_sensor_token': 'r5491t78vlex3qii8gyh3vjp0avkrj47',
226 | 'timestamp': 1696454482897062,
227 | 'fileformat': 'jpg',
228 | 'is_key_frame': True,
229 | 'height': 464,
230 | 'width': 720,
231 | 'filename': 'sweeps/CAM_FRONT_CENTER/1696454482897062.jpg',
232 | 'prev': '',
233 | 'next': '33r4265w297khyvqe033sl2r6m5iylcr',
234 | 'sensor_modality': 'camera',
235 | 'channel': 'CAM_FRONT_CENTER'}
236 | ```
237 | - `ego_pose_token`: token of vehicle ego pose at the time of this sample.
238 | - `calibrated_sensor_token`: token of sensor calibration information (e.g. distortion coefficient, camera intrinsics, sensor pose & location relative to vehicle, etc.).
239 | - `is_key_frame`: disregard; all images have been marked as key frame in our dataset.
240 | - `height`: image height in pixel
241 | - `width`: image width in pixel
242 | - `filename`: image directory relative to the dataset's root folder
243 | - `prev`: previous data token for this sensor
244 | - `next`: next data token for this sensor
245 |
246 |
247 |
248 | After getting the `sample_data` dictionary, Use NuScenes devkit's `get_sample_data()` function to retrieve the data's absolute path.
249 |
250 | Then you may now load the image in any ways you'd like. Here's an example using `cv2`:
251 | ```
252 | import cv2
253 |
254 | data_path, boxes, camera_intrinsic = nusc.get_sample_data(sample_data_token)
255 | img = cv2.imread(data_path)
256 | cv2.imshow('fc_img', img)
257 | cv2.waitKey()
258 | ```
259 |
260 | Output:
261 | ```
262 | ('{$dataset_root}/MARS_multitraversal/10/sweeps/CAM_FRONT_CENTER/1696454482897062.jpg',
263 | [],
264 | array([[661.094568 , 0. , 370.6625195],
265 | [ 0. , 657.7004865, 209.509716 ],
266 | [ 0. , 0. , 1. ]]))
267 | ```
268 | 
269 |
270 | ---
271 | ### LiDAR Data
272 | Same as loading camera data, we start with querying the `sample_data` dictionary for LiDAR sensor.
273 |
274 | Impoirt data calss "LidarPointCloud" from NuScenes devkit for convenient lidar pcd loading and manipulation.
275 |
276 | The `.bcd.bin` LiDAR data in our dataset has 5 dimensions: [ x || y || z || intensity || ring ].
277 |
278 | The 5-dimensional data array is in `pcd.points`. Below is an example of visualizing the pcd with Open3d interactive visualizer.
279 |
280 |
281 | ```
282 | import open3d as o3d
283 | from nuscenes.utils.data_classes import LidarPointCloud
284 |
285 | sensor = 'LIDAR_FRONT_CENTER'
286 | sample_data_token = my_sample['data'][sensor]
287 | lidar_data = nusc.get('sample_data', sample_data_token)
288 |
289 | data_path, boxes, _ = nusc.get_sample_data(my_sample['data'][sensor])
290 |
291 | pcd = LidarPointCloud.from_file(data_path)
292 | print(pcd.points)
293 | pts = pcd.points[:3].T
294 |
295 | # open3d visualizer
296 | vis1 = o3d.visualization.Visualizer()
297 | vis1.create_window(
298 | window_name='pcd viewer',
299 | width=256 * 4,
300 | height=256 * 4,
301 | left=480,
302 | top=270)
303 | vis1.get_render_option().background_color = [0, 0, 0]
304 | vis1.get_render_option().point_size = 1
305 | vis1.get_render_option().show_coordinate_frame = True
306 |
307 | o3d_pcd = o3d.geometry.PointCloud()
308 | o3d_pcd.points = o3d.utility.Vector3dVector(pts)
309 |
310 | vis1.add_geometry(o3d_pcd)
311 | while True:
312 | vis1.update_geometry(o3d_pcd)
313 | vis1.poll_events()
314 | vis1.update_renderer()
315 | time.sleep(0.005)
316 | ```
317 |
318 | Output:
319 | ```
320 | 5-d lidar data:
321 | [[ 3.7755847e+00 5.0539265e+00 5.4277039e+00 ... 3.1050100e+00
322 | 3.4012783e+00 3.7089713e+00]
323 | [-6.3800979e+00 -7.9569578e+00 -7.9752398e+00 ... -7.9960880e+00
324 | -7.9981585e+00 -8.0107889e+00]
325 | [-1.5409404e+00 -3.2752687e-01 5.7313687e-01 ... 5.5921113e-01
326 | -7.5427920e-01 6.6252775e-02]
327 | [ 9.0000000e+00 1.6000000e+01 1.4000000e+01 ... 1.1000000e+01
328 | 1.8000000e+01 1.6000000e+01]
329 | [ 4.0000000e+00 5.3000000e+01 1.0200000e+02 ... 1.0500000e+02
330 | 2.6000000e+01 7.5000000e+01]]
331 | ```
332 |
333 | 
334 |
335 |
336 | ---
337 | ### IMU Data
338 | IMU data in our dataset is saved as json files.
339 | ```
340 | sensor = 'IMU_TOP'
341 | sample_data_token = my_sample['data'][sensor]
342 | lidar_data = nusc.get('sample_data', sample_data_token)
343 |
344 | data_path, boxes, _ = nusc.get_sample_data(my_sample['data'][sensor])
345 |
346 | imu_data = json.load(open(data_path))
347 | print(imu_data)
348 | ```
349 |
350 | Output:
351 | ```
352 | {'utime': 1696454482879084,
353 | 'lat': 42.28098291158676,
354 | 'lon': -83.74725341796875,
355 | 'elev': 259.40500593185425,
356 | 'vel': [0.19750464521348476, -4.99952995654127e-27, -0.00017731071625348704],
357 | 'avel': [-0.0007668623868539726, -0.0006575787383553688, 0.0007131154834496556],
358 | 'acc': [-0.28270150907337666, -0.03748669268679805, 9.785771369934082]}
359 | ```
360 | - `lat`: GPS latitude.
361 | - `lon`: GPS longitude.
362 | - `elev`: GPS elevation.
363 | - `vel`: vehicle instant velocity [x, y, z] in m/s.
364 | - `avel`: vehicle instant angular velocity [x, y, z] in rad/s.
365 | - `acc`: vehicle instant acceleration [x, y, z] in m/s^2.
366 |
367 | ---
368 | ### Vehicle and Sensor Pose
369 | Poses are represented as one rotation matrix and one translation matrix.
370 | - rotation: quaternion [w, x, y, z]
371 | - translation: [x, y, z] in meters
372 |
373 | Sensor-to-vehicle poses may differ for different vehicles. But for each vehicle, its sensor poses should remain unchanged across all scenes & samples.
374 |
375 | Vehicle ego pose can be quaried from sensor data. It should be the same for all sensors in the same sample.
376 |
377 | ```
378 | # get the vehicle ego pose at the time of this FC_data
379 | vehicle_pose_fc = nusc.get('ego_pose', FC_data['ego_pose_token'])
380 | print("vehicle pose: \n", vehicle_pose_fc, "\n")
381 |
382 | # get the vehicle ego pose at the time of this lidar_data, should be the same as that queried from FC_data as they are from the same sample.
383 | vehicle_pose = nusc.get('ego_pose', lidar_data['ego_pose_token'])
384 | print("vehicle pose: \n", vehicle_pose, "\n")
385 |
386 | # get camera pose relative to vehicle at the time of this sample
387 | fc_pose = nusc.get('calibrated_sensor', FC_data['calibrated_sensor_token'])
388 | print("CAM_FRONT_CENTER pose: \n", fc_pose, "\n")
389 |
390 | # get lidar pose relative to vehicle at the time of this sample
391 | lidar_pose = nusc.get('calibrated_sensor', lidar_data['calibrated_sensor_token'])
392 | print("CAM_FRONT_CENTER pose: \n", lidar_pose)
393 | ```
394 |
395 | Output:
396 | ```
397 | vehicle pose:
398 | {'token': 'q9e0pgk3wiot983g4ha8178zrnr37m50',
399 | 'timestamp': 1696454482883182,
400 | 'rotation': [-0.7174290249840286, 0.0, -0.0, -0.6966316057361065],
401 | 'translation': [-146.83352790433003, -21.327001411798392, 0.0]}
402 |
403 | vehicle pose:
404 | {'token': '13y90okaf208cqqy1v54z87cpv88k2qy',
405 | 'timestamp': 1696454482883182,
406 | 'rotation': [-0.7174290249840286, 0.0, -0.0, -0.6966316057361065],
407 | 'translation': [-146.83352790433003, -21.327001411798392, 0.0]}
408 |
409 | CAM_FRONT_CENTER pose:
410 | {'token': 'r5491t78vlex3qii8gyh3vjp0avkrj47',
411 | 'sensor_token': '1gk062vf442xsn86xo152qw92596k8b9',
412 | 'translation': [2.24715, 0.0, 1.4725],
413 | 'rotation': [0.49834929780875276, -0.4844970241435727, 0.5050790448056688, -0.5116695901338464],
414 | 'camera_intrinsic': [[661.094568, 0.0, 370.6625195], [0.0, 657.7004865, 209.509716], [0.0, 0.0, 1.0]],
415 | 'distortion_coefficient': [0.122235, -1.055498, 2.795589, -2.639154]}
416 |
417 | CAM_FRONT_CENTER pose:
418 | {'token': '6f367iy1b5c97e8gu614n63jg1f5os19',
419 | 'sensor_token': 'myfmnd47g91ijn0a7481eymfk253iwy9',
420 | 'translation': [2.12778, 0.0, 1.57],
421 | 'rotation': [0.9997984797097376, 0.009068089160690487, 0.006271772522201215, -0.016776012592418482]}
422 |
423 | ```
424 |
425 |
426 |
427 | ## LiDAR-Image projection
428 | - Use NuScenes devkit's `render_pointcloud_in_image()` method.
429 | - The first variable is a sample token.
430 | - Use `camera_channel` to specify the camera name you'd like to project the poiint cloud onto.
431 | ```
432 | nusc.render_pointcloud_in_image(my_sample['token'],
433 | pointsensor_channel='LIDAR_FRONT_CENTER',
434 | camera_channel='CAM_FRONT_CENTER',
435 | render_intensity=False,
436 | show_lidarseg=False)
437 | ```
438 |
439 | Output:
440 |
441 | 
442 |
443 |
444 |
445 |
446 |
447 |
448 |
449 |
--------------------------------------------------------------------------------
/visualize/lidar_cam_viz_agent.py:
--------------------------------------------------------------------------------
1 | import os.path
2 | import time
3 |
4 | import matplotlib
5 | from nuscenes import NuScenes
6 | from nuscenes.utils.geometry_utils import BoxVisibility, transform_matrix, view_points
7 | from nuscenes.utils.data_classes import LidarPointCloud, RadarPointCloud, Box
8 | import numpy as np
9 | import open3d as o3d
10 | from pyquaternion import Quaternion
11 | import cv2
12 | from tqdm import tqdm
13 | import math
14 | import datetime
15 |
16 | img_width = 720
17 | img_height = 464
18 | img_divider = 1
19 | frame_width = 4
20 | font = cv2.FONT_HERSHEY_DUPLEX
21 | color_dict = {"maisy": (200, 50, 0), "metheven": (0, 125, 250), "mike": (50, 150, 0), "mithrandir": (150, 0, 150)}
22 | # color_dict = {"maisy": (200, 0, 0), "marinara": (0, 200, 0), "metheven": (0, 0, 200), "mike": (200, 200, 0), "mithrandir": (200, 0, 200)}
23 | # color_dict = {"maisy": (100, 100, 0), "metheven": (200, 200, 20), "mike": (50, 200, 200), "mithrandir": (120, 0, 120)}
24 | # color_dict = {"maisy": (0, 80, 250), "metheven": (0, 190, 80), "mike": (200, 200, 0), "mithrandir": (200, 0, 200)}
25 | color_idx_dict = {"maisy": 0, "marinara": 1, "metheven": 2, "mike": 3, "mithrandir": 4}
26 | v_id_dict = {"maisy": "Agent 0", "marinara": "Agent 1", "metheven": "Agent 2", "mike": "Agent 3", "mithrandir": "Agent 4"}
27 | cmap = matplotlib.colormaps['Set1']
28 |
29 | def transform_pcd(pcd, tran_matrix):
30 | '''
31 | :param pcd: 3*n np array (1st row x, 2nd row y, 3rd row z)
32 | :param tran_matrix: 4*4 transformation matrix
33 | :return: n*3 pcd (each row is a point [x, y, z])
34 | '''
35 | # append the pt by 1 to apply the transformation matrix
36 | rot = tran_matrix[:3, :3]
37 | trl = tran_matrix[:3, 3]
38 | pcd = (rot @ pcd).T + trl
39 |
40 | return pcd
41 |
42 |
43 | def get_scene_tokens(nusc, first_sample_token):
44 | token_list = []
45 | sample = nusc.get('sample', first_sample_token)
46 | scene = nusc.get('scene', sample["scene_token"])
47 |
48 | nbr_samples = scene['nbr_samples']
49 | last_sample_token = scene['last_sample_token']
50 | token_list.append(first_sample_token)
51 | print(scene["name"])
52 | # vehicle = scene["name"].split("_")[-1]
53 | for i in range(nbr_samples - 1):
54 | if sample == last_sample_token:
55 | break
56 | token_list.append(sample['next'])
57 | sample = nusc.get('sample', sample['next'])
58 | # return token_list, vehicle
59 | return token_list
60 |
61 |
62 | def get_scene_list(nusc):
63 | scene_dict = {}
64 | for i in range(len(nusc.scene)):
65 | scene = nusc.scene[i]
66 | first_sample_token = scene['first_sample_token']
67 | vehicle = scene["vehicle_list"]
68 | sample = nusc.get('sample', first_sample_token)
69 | tp = int(sample["timestamp"])
70 | scene_dict[first_sample_token] = [tp, vehicle]
71 | scene_dict = dict(sorted(scene_dict.items(), key=lambda item: item[1][0]))
72 | # print(scene_dict)
73 |
74 | return scene_dict
75 |
76 |
77 | def img_frame(img, color):
78 |
79 | # print(color.astype(np.uint8))
80 | framed_img = cv2.copyMakeBorder(
81 | img,
82 | top=frame_width,
83 | bottom=frame_width,
84 | left=frame_width,
85 | right=frame_width,
86 | borderType=cv2.BORDER_CONSTANT,
87 | value=color[0].tolist()
88 | )
89 |
90 | # cv2.imshow("framed cam", framed_img)
91 | # cv2.waitKey()
92 | return framed_img
93 |
94 |
95 | def add_tag(img, tp, vehicle, color=(0, 0, 0), fontScale=1.2, thickness=2):
96 | v_color = color_dict[vehicle]
97 |
98 | date_time = datetime.datetime.fromtimestamp(int(tp)/1000000).strftime('%Y-%m-%d %H:%M:%S.%f')
99 | date = date_time.split(" ")[0]
100 | time_s = date_time.split(" ")[1][:-3]
101 | h = img.shape[0]
102 | w = img.shape[1]
103 | # blank = np.ones((h, h, 3)) * 255
104 | blank = cv2.imread("calender_img_edge.png")
105 | blank = cv2.resize(blank, dsize=(h, h), interpolation=cv2.INTER_CUBIC)
106 | org1 = (h//8+50, (h//2))
107 | time_tag = cv2.putText(blank, date, org1, font, fontScale, color, thickness, cv2.LINE_AA)
108 | org2 = (h//8+50, (h//2)+80)
109 | time_tag = cv2.putText(time_tag, time_s, org2, font, fontScale, color, thickness, cv2.LINE_AA)
110 | org3 = (h//8+98, (h//2)+160)
111 | time_tag = cv2.putText(time_tag, v_id_dict[vehicle], org3, font, fontScale, v_color, thickness, cv2.LINE_AA)
112 | result = np.concatenate((time_tag, img), axis=1)
113 | return result
114 |
115 |
116 | if __name__ == '__main__':
117 | root_dir = f"I:/MARS_agent_10Hz_40ms"
118 | nusc = NuScenes(version='v1.0', dataroot=root_dir, verbose=True)
119 | scene_dict = get_scene_list(nusc)
120 | nbr_scene = len(nusc.scene)
121 | # nbr_scene = 10
122 | # step = len(nusc.scene)//nbr_scene
123 | # print(f"Multi Agent has {len(nusc.scene)} scenes, showing first {nbr_scene}")
124 |
125 | ''' idx of scene you'd like to visualize '''
126 | # scene 21 has three cars, scene 30 has encounter from opposite direction
127 | scene_idx = 21
128 |
129 | scene = nusc.scene[scene_idx]
130 | first_token_list = list(scene_dict.keys())
131 | first_token = first_token_list[scene_idx]
132 | vehicles = scene_dict[first_token][1]
133 | all_token_list = get_scene_tokens(nusc, first_token)
134 | nbr_sample = len(all_token_list)
135 | nbr_vehicle = len(vehicles)
136 |
137 | cap = cv2.VideoCapture(0)
138 | fourcc = cv2.VideoWriter_fourcc(*'MJPG')
139 | ''' camera + lidar width + frame width + time tag width '''
140 | out0 = cv2.VideoWriter(f"{root_dir}/agent_80m_{scene_idx}.avi", fourcc, 30,
141 | ((img_width * 3 + img_height * nbr_vehicle//3*2)//img_divider + frame_width*2 + img_height//img_divider -4,
142 | (img_height//img_divider + frame_width*2) * nbr_vehicle))
143 |
144 | vis1 = o3d.visualization.Visualizer()
145 | vis1.create_window(
146 | window_name='viz',
147 | width=img_height // img_divider //3*2 * nbr_vehicle,
148 | height=(img_height // img_divider + frame_width*2) * nbr_vehicle,
149 | left=480,
150 | top=270)
151 | vis1.get_render_option().background_color = [1, 1, 1]
152 | vis1.get_render_option().point_size = 1
153 | vis1.get_render_option().show_coordinate_frame = True
154 |
155 | # empty canvas for img display
156 | img_canvas = np.ones(((img_height // img_divider + frame_width*2)*nbr_vehicle, img_width // img_divider * 3 + frame_width*2 + img_height // img_divider, 3)) * 255
157 | # empty point cloud
158 | pcd_obj = o3d.geometry.PointCloud()
159 |
160 | start_flag = True
161 | for i in tqdm(range(1, nbr_sample), desc=f"scene {scene_idx} of {nbr_scene}"):
162 |
163 | if start_flag:
164 | sample = nusc.get('sample', first_token)
165 | else:
166 | if sample['next'] == '':
167 | break
168 | else:
169 | sample = nusc.get('sample', sample['next'])
170 |
171 | sample_pcd_list = []
172 | sample_color_list = []
173 | # print(vehicles)
174 | for j in range(len(vehicles)):
175 | vehicle = vehicles[j]
176 | ''' ======================================== lidar ======================================== '''
177 | sample_lidar_token = sample['data'][f'LIDAR_FRONT_CENTER_{vehicle}']
178 | lidar_data = nusc.get('sample_data', sample_lidar_token)
179 |
180 | pcd_path = f"{nusc.dataroot}/{lidar_data['filename']}"
181 | pcd = LidarPointCloud.from_file(pcd_path).points
182 |
183 | ''' preserve only pts within 40 meter radius '''
184 | mask = np.sqrt((pcd[0, :] ** 2 + pcd[1, :] ** 2)) < 80
185 | pcd = pcd[:, mask][:3, :]
186 | # pcd = pcd[:3, :]
187 |
188 | ''' transform pcd from local frame to global frame '''
189 | ego_pose = nusc.get('ego_pose', lidar_data['ego_pose_token'])
190 |
191 | lidar_pose = nusc.get('calibrated_sensor', lidar_data['calibrated_sensor_token'])
192 | # Get global pose
193 | ego_pose = transform_matrix(ego_pose['translation'], Quaternion(ego_pose['rotation']), inverse=False)
194 | lidar_pose = transform_matrix(lidar_pose['translation'], Quaternion(lidar_pose['rotation']), inverse=False)
195 | global_pose = np.dot(ego_pose, lidar_pose)
196 |
197 | pcd = transform_pcd(pcd, global_pose)
198 |
199 | # color = np.array(color_dict[vehicle])
200 | color = np.array([color_dict[vehicle]])/255
201 | # print(vehicle, color)
202 | # color = np.array([cmap(1 / 5 * color_idx_dict[vehicle]*2)[:3]])
203 | colors = np.repeat(color, len(pcd), axis=0)
204 | sample_pcd_list.append(pcd)
205 | sample_color_list.append(colors)
206 |
207 | ''' ======================================== camera ======================================== '''
208 | sample_fc_token = sample['data'][f'CAM_FRONT_CENTER_{vehicle}']
209 | sample_fl_token = sample['data'][f'CAM_FRONT_LEFT_{vehicle}']
210 | sample_fr_token = sample['data'][f'CAM_FRONT_RIGHT_{vehicle}']
211 |
212 | fc_data = nusc.get('sample_data', sample_fc_token)
213 | fc_path = f"{nusc.dataroot}/{fc_data['filename']}"
214 | fc_img = cv2.imread(fc_path)
215 |
216 | fl_data = nusc.get('sample_data', sample_fl_token)
217 | fl_path = f"{nusc.dataroot}/{fl_data['filename']}"
218 | fl_img = cv2.imread(fl_path)
219 |
220 | fr_data = nusc.get('sample_data', sample_fr_token)
221 | fr_path = f"{nusc.dataroot}/{fr_data['filename']}"
222 | fr_img = cv2.imread(fr_path)
223 |
224 | h = max([fc_img.shape[0], fl_img.shape[0], fr_img.shape[0]])
225 | w = max([fc_img.shape[1], fl_img.shape[1], fr_img.shape[1]])
226 |
227 | fc_img = cv2.resize(fc_img, dsize=(w // img_divider, h // img_divider), interpolation=cv2.INTER_CUBIC)
228 | fl_img = cv2.resize(fl_img, dsize=(w // img_divider, h // img_divider), interpolation=cv2.INTER_CUBIC)
229 | fr_img = cv2.resize(fr_img, dsize=(w // img_divider, h // img_divider), interpolation=cv2.INTER_CUBIC)
230 |
231 | img_arr1 = [fl_img, fc_img, fr_img]
232 |
233 | vehicle_img = np.concatenate(img_arr1, axis=1)
234 | # vehicle_img2 = np.concatenate(img_arr2, axis=1)
235 | # vehicle_img = np.concatenate((vehicle_img1, vehicle_img2), axis=0)
236 | tp = lidar_data['filename'].split("/")[-1].split(".")[0]
237 |
238 | vehicle_img = add_tag(vehicle_img, tp, vehicle)
239 | # vehicle_img = img_frame(vehicle_img, color)
240 |
241 | num_row = vehicle_img.shape[0]
242 | num_col = vehicle_img.shape[1]
243 | img_canvas[num_row * j:num_row * (j + 1), :num_col, :] = vehicle_img
244 | img_canvas = img_canvas.astype(np.uint8)
245 |
246 | ''' new lidar scan '''
247 | pcd_obj.points = o3d.utility.Vector3dVector(np.concatenate(sample_pcd_list, axis=0))
248 | pcd_obj.colors = o3d.utility.Vector3dVector(np.concatenate(sample_color_list, axis=0))
249 |
250 |
251 | if start_flag:
252 | vis1.add_geometry(pcd_obj)
253 | start_flag = False
254 | view = vis1.get_view_control()
255 | view.set_zoom(0.5)
256 | if scene_idx == 1:
257 | view.translate(-50, -100)
258 | elif scene_idx == 2:
259 | view.translate(-50, -100)
260 | elif scene_idx == 5:
261 | view.translate(0, 60)
262 | elif scene_idx == 18:
263 | view.translate(50, -300)
264 | elif scene_idx == 21:
265 | view.translate(-150, 0)
266 | elif scene_idx == 29:
267 | view.translate(150, -0)
268 | elif scene_idx == 30:
269 | view.translate(0, -50)
270 |
271 |
272 | else:
273 | vis1.update_geometry(pcd_obj)
274 |
275 | vis1.poll_events()
276 | vis1.update_renderer()
277 | # cv2.imshow("cam", img_canvas)
278 | # cv2.waitKey(100)
279 | time.sleep(0.02)
280 |
281 | ''' write into video '''
282 | o3d_screenshot_1 = vis1.capture_screen_float_buffer()
283 | o3d_screenshot_1 = (255.0 * np.asarray(o3d_screenshot_1)).astype(np.uint8)
284 | out0.write(np.concatenate((img_canvas, o3d_screenshot_1), axis=1))
285 |
286 | ''' save some screenshots '''
287 | # if i % 5 == 0:
288 | # img_out_dir = f"{root_dir}/img/{scene_idx}"
289 | # if not os.path.exists(img_out_dir):
290 | # os.makedirs(img_out_dir)
291 | # cv2.imwrite(f"{img_out_dir}/{i}.jpg", np.concatenate((img_canvas, o3d_screenshot_1), axis=1))
292 | # cv2.imwrite(f"{img_out_dir}/{i}_lidar.jpg", o3d_screenshot_1)
293 |
294 |
--------------------------------------------------------------------------------
/visualize/lidar_cam_viz_traversal.py:
--------------------------------------------------------------------------------
1 | import os.path
2 | import time
3 |
4 | import matplotlib
5 | from nuscenes import NuScenes
6 | from nuscenes.utils.geometry_utils import BoxVisibility, transform_matrix, view_points
7 | from nuscenes.utils.data_classes import LidarPointCloud, RadarPointCloud, Box
8 | import numpy as np
9 | import open3d as o3d
10 | from pyquaternion import Quaternion
11 | import cv2
12 | from tqdm import tqdm
13 | import math
14 | import datetime
15 |
16 | img_width = 720
17 | img_height = 464
18 | img_divider = 4
19 | frame_width = 4
20 | font = cv2.FONT_HERSHEY_DUPLEX
21 |
22 | color_dict = {"maisy": (200, 0, 0), "marinara": (0, 200, 0), "metheven": (0, 0, 200), "mike": (200, 200, 0), "mithrandir": (200, 0, 200)}
23 | cmap = matplotlib.colormaps['tab20']
24 |
25 | def transform_pcd(pcd, tran_matrix):
26 | '''
27 | :param pcd: 3*n np array (1st row x, 2nd row y, 3rd row z)
28 | :param tran_matrix: 4*4 transformation matrix
29 | :return: n*3 pcd (each row is a point [x, y, z])
30 | '''
31 | # append the pt by 1 to apply the transformation matrix
32 | rot = tran_matrix[:3, :3]
33 | trl = tran_matrix[:3, 3]
34 | pcd = (rot @ pcd).T + trl
35 |
36 | return pcd
37 |
38 |
39 | def get_scene_tokens(nusc, first_sample_token):
40 | token_list = []
41 | sample = nusc.get('sample', first_sample_token)
42 | scene = nusc.get('scene', sample["scene_token"])
43 |
44 | nbr_samples = scene['nbr_samples']
45 | last_sample_token = scene['last_sample_token']
46 | token_list.append(first_sample_token)
47 | print(scene["name"])
48 | # vehicle = scene["name"].split("_")[-1]
49 | for i in range(nbr_samples - 1):
50 | if sample == last_sample_token:
51 | break
52 | token_list.append(sample['next'])
53 | sample = nusc.get('sample', sample['next'])
54 | # return token_list, vehicle
55 | return token_list
56 |
57 |
58 | def get_scene_list(nusc):
59 | scene_dict = {}
60 | for i in range(len(nusc.scene)):
61 | scene = nusc.scene[i]
62 | first_sample_token = scene['first_sample_token']
63 | vehicle = scene["name"].split("_")[-1]
64 | sample = nusc.get('sample', first_sample_token)
65 | tp = int(sample["timestamp"])
66 | scene_dict[first_sample_token] = [tp, vehicle]
67 | scene_dict = dict(sorted(scene_dict.items(), key=lambda item: item[1][0]))
68 | print(scene_dict)
69 |
70 | return scene_dict
71 |
72 |
73 | def img_frame(img, color):
74 | framed_img = cv2.copyMakeBorder(
75 | img,
76 | top=frame_width,
77 | bottom=frame_width,
78 | left=frame_width,
79 | right=frame_width,
80 | borderType=cv2.BORDER_CONSTANT,
81 | value=color[0]*255
82 | )
83 | # print(color[0]*255)
84 |
85 | # cv2.imshow("framed cam", framed_img)
86 | # cv2.waitKey()
87 | return framed_img
88 |
89 |
90 | def add_tag(img, tp, vehicle, color=(0, 0, 0), fontScale=0.4, thickness=1):
91 | date_time = datetime.datetime.fromtimestamp(int(tp)/1000000).strftime('%Y-%m-%d %H:%M:%S.%f')
92 | date = date_time.split(" ")[0]
93 | time_s = date_time.split(" ")[1][:-3]
94 | h = img.shape[0]
95 | w = img.shape[1]
96 | # blank = np.ones((h, h, 3)) * 255
97 | blank = cv2.imread("calender_img_edge.png")
98 | blank = cv2.resize(blank, dsize=(h, h), interpolation=cv2.INTER_CUBIC)
99 | org1 = (h//8, (h//2))
100 | time_tag = cv2.putText(blank, date, org1, font, fontScale, color, thickness, cv2.LINE_AA)
101 | org2 = (h//8, (h//2)+20)
102 | time_tag = cv2.putText(time_tag, time_s, org2, font, fontScale, color, thickness, cv2.LINE_AA)
103 | org3 = (h//8, (h//2)+38)
104 |
105 | time_tag = cv2.putText(time_tag, vehicle, org3, font, fontScale, color_dict[vehicle], thickness, cv2.LINE_AA)
106 | result = np.concatenate((time_tag, img), axis=1)
107 | return result
108 |
109 |
110 | if __name__ == '__main__':
111 | root_dir = "K:/MARS_10Hz_location"
112 | location_list = [0, 1, 2, 3, 4, 5]
113 |
114 | for location in location_list:
115 | if str(location) not in os.listdir(root_dir):
116 | continue
117 |
118 | root_dir = f"K:/MARS_10Hz_location/{location}"
119 | nusc = NuScenes(version='v1.0', dataroot=root_dir, verbose=True)
120 | scene_dict = get_scene_list(nusc)
121 |
122 | ''' number of scenes to be visualized '''
123 | # nbr_scene = len(nusc.scene)
124 | nbr_scene = 10
125 | step = len(nusc.scene)//nbr_scene
126 | print(f"location {location} has {len(nusc.scene)} scenes, showing first {nbr_scene}")
127 |
128 | cap = cv2.VideoCapture(0)
129 | fourcc = cv2.VideoWriter_fourcc(*'MJPG')
130 | out0 = cv2.VideoWriter(f"{root_dir}/../{location}_80m.avi", fourcc, 120,
131 | ((img_width * 3 + img_height * 10)//img_divider + frame_width*2 + img_height//img_divider,
132 | (img_height//img_divider + frame_width*2) * 10))
133 |
134 | vis1 = o3d.visualization.Visualizer()
135 | vis1.create_window(
136 | window_name='viz',
137 | width=img_height // img_divider * 10,
138 | height=(img_height // img_divider + frame_width * 2) * 10,
139 | left=480,
140 | top=270)
141 | vis1.get_render_option().background_color = [1, 1, 1]
142 | vis1.get_render_option().point_size = 1
143 | vis1.get_render_option().show_coordinate_frame = True
144 |
145 |
146 | ''' empty canvas for img display '''
147 | img_canvas = np.ones(((img_height // img_divider + frame_width*2)*10, img_width // img_divider * 3 + frame_width*2 + img_height // img_divider, 3)) * 255
148 |
149 | ''' open3d point cloud object '''
150 | pcd_obj = o3d.geometry.PointCloud()
151 |
152 | sample_token_dict = {}
153 | frame_offset_dict = {}
154 |
155 | start_flag = True
156 | first_token_list = list(scene_dict.keys())
157 | for i in range(nbr_scene):
158 | first_token = first_token_list[i*step]
159 | vehicle = scene_dict[first_token][1]
160 | sample_token_dict[str(i)] = [vehicle, get_scene_tokens(nusc, first_token)]
161 | frame_offset_dict[str(i)] = 0
162 | nbr_sample = len(sample_token_dict[str(i)])
163 |
164 | frame = 0
165 | finish_count = 0
166 | skip_idx = []
167 | batch = [0]
168 |
169 | last_scene_idx = max(batch)
170 |
171 | while True:
172 | location_pcd_list = []
173 | location_color_list = []
174 | img_list = []
175 |
176 | for idx in range(len(batch)):
177 | i = batch[idx]
178 | ''' if scene_i is finished playing, skip it '''
179 | if i in skip_idx:
180 | continue
181 |
182 | scene_idx = str(i)
183 | offset = frame_offset_dict[scene_idx]
184 | ''' if last frame (frame - offset >= nbr samples) '''
185 | if frame - offset >= len(sample_token_dict[scene_idx][1]):
186 | finish_count += 1
187 | skip_idx.append(i)
188 |
189 | last_scene_idx += 1
190 | # if last_scene_idx < nbr_scene:
191 | if last_scene_idx < nbr_scene:
192 | ''' substitute the finished scene_idx with a new one '''
193 | # if len(batch) == 1:
194 | target_idx = batch.index(i)
195 | batch[target_idx] = last_scene_idx
196 | frame_offset_dict[str(last_scene_idx)] = frame
197 | print(f"{frame} scene {i} finished, showing {batch}")
198 | else:
199 | print(f"{frame} scene {i} finished, this is the last scene")
200 | continue
201 |
202 | sample_token = sample_token_dict[scene_idx][1][frame - offset]
203 | sample = nusc.get('sample', sample_token)
204 | # print(sample['data'])
205 |
206 | ''' ======================================== lidar ======================================== '''
207 | sample_lidar_token = sample['data']['LIDAR_FRONT_CENTER']
208 | lidar_data = nusc.get('sample_data', sample_lidar_token)
209 | pcd_path = f"{nusc.dataroot}/{lidar_data['filename']}"
210 | pcd = LidarPointCloud.from_file(pcd_path).points
211 |
212 | ''' preserve only pts within 40 meter radius '''
213 | mask = np.sqrt((pcd[0, :] ** 2 + pcd[1, :] ** 2)) < 80
214 | pcd = pcd[:, mask][:3, :]
215 | # pcd = pcd[:3, :]
216 |
217 | ''' transform pcd from local frame to global frame '''
218 | ego_pose = nusc.get('ego_pose', lidar_data['ego_pose_token'])
219 | lidar_pose = nusc.get('calibrated_sensor', lidar_data['calibrated_sensor_token'])
220 | # Get global pose
221 | ego_pose = transform_matrix(ego_pose['translation'], Quaternion(ego_pose['rotation']), inverse=False)
222 | lidar_pose = transform_matrix(lidar_pose['translation'], Quaternion(lidar_pose['rotation']), inverse=False)
223 | global_pose = np.dot(ego_pose, lidar_pose)
224 |
225 | pcd = transform_pcd(pcd, global_pose)
226 | # color = np.array([color_list[idx]]) / 255
227 | color = np.array([cmap(1/10 * i)[:3]])
228 | colors = np.repeat(color, len(pcd), axis=0)
229 |
230 | location_pcd_list.append(pcd)
231 | location_color_list.append(colors)
232 |
233 | ''' ======================================== camera ======================================== '''
234 | ''' visualize only the front three cameras '''
235 | sample_fc_token = sample['data']['CAM_FRONT_CENTER']
236 | sample_fl_token = sample['data']['CAM_FRONT_LEFT']
237 | sample_fr_token = sample['data']['CAM_FRONT_RIGHT']
238 | # sample_bc_token = sample['data']['CAM_BACK_CENTER']
239 | # sample_lc_token = sample['data']['CAM_SIDE_LEFT']
240 | # sample_rc_token = sample['data']['CAM_SIDE_RIGHT']
241 |
242 | fc_data = nusc.get('sample_data', sample_fc_token)
243 | fc_path = f"{nusc.dataroot}/{fc_data['filename']}"
244 | fc_img = cv2.imread(fc_path)
245 |
246 | fl_data = nusc.get('sample_data', sample_fl_token)
247 | fl_path = f"{nusc.dataroot}/{fl_data['filename']}"
248 | fl_img = cv2.imread(fl_path)
249 |
250 | fr_data = nusc.get('sample_data', sample_fr_token)
251 | fr_path = f"{nusc.dataroot}/{fr_data['filename']}"
252 | fr_img = cv2.imread(fr_path)
253 |
254 | # bc_data = nusc.get('sample_data', sample_bc_token)
255 | # bc_path = f"{nusc.dataroot}/{bc_data['filename']}"
256 | # bc_img = cv2.imread(bc_path)
257 | #
258 | # lc_data = nusc.get('sample_data', sample_lc_token)
259 | # lc_path = f"{nusc.dataroot}/{lc_data['filename']}"
260 | # lc_img = cv2.imread(lc_path)
261 | #
262 | # rc_data = nusc.get('sample_data', sample_rc_token)
263 | # rc_path = f"{nusc.dataroot}/{rc_data['filename']}"
264 | # rc_img = cv2.imread(rc_path)
265 |
266 | h = max([fc_img.shape[0], fl_img.shape[0], fr_img.shape[0]])
267 | w = max([fc_img.shape[1], fl_img.shape[1], fr_img.shape[1]])
268 |
269 | fc_img = cv2.resize(fc_img, dsize=(w//img_divider, h//img_divider), interpolation=cv2.INTER_CUBIC)
270 | fl_img = cv2.resize(fl_img, dsize=(w//img_divider, h//img_divider), interpolation=cv2.INTER_CUBIC)
271 | fr_img = cv2.resize(fr_img, dsize=(w//img_divider, h//img_divider), interpolation=cv2.INTER_CUBIC)
272 | # bc_img = cv2.resize(bc_img, dsize=(w, h), interpolation=cv2.INTER_CUBIC)
273 | # lc_img = cv2.resize(lc_img, dsize=(w, h), interpolation=cv2.INTER_CUBIC)
274 | # rc_img = cv2.resize(rc_img, dsize=(w, h), interpolation=cv2.INTER_CUBIC)
275 |
276 | img_arr1 = [fl_img, fc_img, fr_img]
277 | # img_arr2 = [lc_img, bc_img, rc_img]
278 |
279 | vehicle_img = np.concatenate(img_arr1, axis=1)
280 | # vehicle_img2 = np.concatenate(img_arr2, axis=1)
281 | # vehicle_img = np.concatenate((vehicle_img1, vehicle_img2), axis=0)
282 | tp = lidar_data['filename'].split("/")[-1].split(".")[0]
283 | vehicle = sample_token_dict[scene_idx][0]
284 | vehicle_img = add_tag(vehicle_img, tp, vehicle)
285 |
286 | vehicle_img = img_frame(vehicle_img, color)
287 |
288 | num_row = vehicle_img.shape[0]
289 | num_col = vehicle_img.shape[1]
290 | img_canvas[num_row*i:num_row*(i+1), :num_col, :] = vehicle_img
291 | img_canvas = img_canvas.astype(np.uint8)
292 |
293 | if finish_count == nbr_scene:
294 | img_out_dir = f"K:/MARS_10Hz_location/img/{location}/"
295 | if not os.path.exists(img_out_dir):
296 | os.makedirs(img_out_dir)
297 | cv2.imwrite(f"{img_out_dir}/{frame}.jpg", np.concatenate((img_canvas, o3d_screenshot_1), axis=1))
298 | vis1.close()
299 | break
300 |
301 | if len(location_pcd_list) == 0:
302 | continue
303 |
304 | ''' new lidar scan '''
305 | pcd_obj.points = o3d.utility.Vector3dVector(np.concatenate(location_pcd_list, axis=0))
306 | pcd_obj.colors = o3d.utility.Vector3dVector(np.concatenate(location_color_list, axis=0))
307 |
308 | ''' new background (down-sampled) '''
309 | if start_flag:
310 | background = pcd_obj.voxel_down_sample(voxel_size=4)
311 | background.points = o3d.utility.Vector3dVector(background.points - np.array([0, 0, 2]))
312 | else:
313 | down_sampled = pcd_obj.voxel_down_sample(voxel_size=4)
314 | background.points = o3d.utility.Vector3dVector(np.concatenate((background.points, down_sampled.points - np.array([0, 0, 10])), axis=0))
315 | background.colors = o3d.utility.Vector3dVector(np.concatenate((background.colors, down_sampled.colors), axis=0))
316 | background = background.voxel_down_sample(voxel_size=1)
317 |
318 | ''' new full frame (background + new scan) '''
319 | pcd_obj.points = o3d.utility.Vector3dVector(np.concatenate((background.points, pcd_obj.points), axis=0))
320 | pcd_obj.colors = o3d.utility.Vector3dVector(np.concatenate((background.colors, pcd_obj.colors), axis=0))
321 |
322 | if start_flag:
323 | vis1.add_geometry(pcd_obj)
324 | start_flag = False
325 | view = vis1.get_view_control()
326 |
327 | ''' adjust visualize initial viewpoint as needed '''
328 | if location == 24:
329 | # # 24
330 | view.set_zoom(1)
331 | view.translate(-150, -100)
332 | # # 15
333 | elif location == 15:
334 | view.set_zoom(1)
335 | view.translate(0, -250)
336 | # # 59
337 | elif location == 59:
338 | view.set_zoom(1)
339 | view.translate(-150, -100)
340 | # # 57
341 | elif location == 57:
342 | view.set_zoom(1)
343 | view.translate(-100, 50)
344 | # # 3
345 | elif location == 3:
346 | view.set_zoom(1)
347 | view.translate(-0, -150)
348 | # 61
349 | elif location == 61:
350 | view.set_zoom(1)
351 | view.translate(200, 50)
352 | else:
353 | vis1.update_geometry(pcd_obj)
354 |
355 | vis1.poll_events()
356 | vis1.update_renderer()
357 | time.sleep(0.005)
358 |
359 | ''' write into video '''
360 | o3d_screenshot_1 = vis1.capture_screen_float_buffer()
361 | o3d_screenshot_1 = (255.0 * np.asarray(o3d_screenshot_1)).astype(np.uint8)
362 | out0.write(np.concatenate((img_canvas, o3d_screenshot_1), axis=1))
363 |
364 | ''' save some snapshots '''
365 | # if frame % 30 == 0:
366 | # img_out_dir = f"K:/MARS_10Hz_location/img/{location}/"
367 | # if not os.path.exists(img_out_dir):
368 | # os.makedirs(img_out_dir)
369 | # cv2.imwrite(f"{img_out_dir}/{frame}.jpg", np.concatenate((img_canvas, o3d_screenshot_1), axis=1))
370 |
371 | frame += 1
372 |
--------------------------------------------------------------------------------
/visualize/lidar_viz.py:
--------------------------------------------------------------------------------
1 | import time
2 |
3 | import matplotlib
4 | from nuscenes import NuScenes
5 | from nuscenes.utils.geometry_utils import BoxVisibility, transform_matrix, view_points
6 | from nuscenes.utils.data_classes import LidarPointCloud, RadarPointCloud, Box
7 | import numpy as np
8 | import open3d as o3d
9 | from pyquaternion import Quaternion
10 | import cv2
11 | from tqdm import tqdm
12 | import math
13 |
14 | img_width = 256
15 | img_height = 256
16 | font = cv2.FONT_HERSHEY_SIMPLEX
17 |
18 | color_list = [(255, 102, 102), (255, 255, 102), (102, 255, 102), (102, 255, 255), (178, 102, 255)]
19 | cmap = matplotlib.colormaps['jet']
20 |
21 | def transform_pcd(pcd, tran_matrix):
22 | '''
23 | :param pcd: 3*n np array (1st row x, 2nd row y, 3rd row z)
24 | :param tran_matrix: 4*4 transformation matrix
25 | :return: n*3 pcd (each row is a point [x, y, z])
26 | '''
27 | # append the pt by 1 to apply the transformation matrix
28 | rot = tran_matrix[:3, :3]
29 | trl = tran_matrix[:3, 3]
30 | pcd = (rot @ pcd).T + trl
31 |
32 | return pcd
33 |
34 |
35 | def get_scene_tokens(nusc, scene_index):
36 | token_list = []
37 | scene = nusc.scene[scene_index]
38 | nbr_samples = scene['nbr_samples']
39 | first_sample_token = scene['first_sample_token']
40 | last_sample_token = scene['last_sample_token']
41 | sample = nusc.get('sample', first_sample_token)
42 | token_list.append(first_sample_token)
43 | print(scene["name"])
44 | # vehicle = scene["name"].split("_")[-1]
45 | for i in range(nbr_samples - 1):
46 | if sample == last_sample_token:
47 | break
48 | token_list.append(sample['next'])
49 | sample = nusc.get('sample', sample['next'])
50 | # return token_list, vehicle
51 | return token_list
52 |
53 |
54 | if __name__ == '__main__':
55 | location = 2
56 | root_dir = f"K:/MARS_10Hz_location/{location}"
57 | nusc = NuScenes(version='v1.0', dataroot=root_dir, verbose=True)
58 | nbr_scene = len(nusc.scene)
59 | print(f"location {location} has {nbr_scene} scenes")
60 |
61 | cap = cv2.VideoCapture(0)
62 | fourcc = cv2.VideoWriter_fourcc(*'MJPG')
63 | out0 = cv2.VideoWriter(f"{root_dir}/lidar_view_traversal.avi", fourcc, 10, (img_width * 4, img_height * 4))
64 |
65 | vis1 = o3d.visualization.Visualizer()
66 | vis1.create_window(
67 | window_name='Ego Vehicle Segmented Scene',
68 | width=img_width * 4,
69 | height=img_height * 4,
70 | left=480,
71 | top=270)
72 | vis1.get_render_option().background_color = [0, 0, 0]
73 | vis1.get_render_option().point_size = 2
74 | vis1.get_render_option().show_coordinate_frame = True
75 |
76 | pcd_obj = o3d.geometry.PointCloud()
77 | '''======================================================================================='''
78 | # start_flag = True
79 | # location_pcd_list = []
80 | # for i in tqdm(range(nbr_scene)):
81 | # scene = nusc.scene[i]
82 | #
83 | # nbr_samples = scene['nbr_samples']
84 | # first_sample_token = scene['first_sample_token']
85 | # last_sample_token = scene['last_sample_token']
86 | # sample = nusc.get('sample', first_sample_token)
87 | #
88 | # while True:
89 | # sample_data_token = sample['data']['LIDAR_FRONT_CENTER']
90 | #
91 | # lidar_data = nusc.get('sample_data', sample_data_token)
92 | # pcd_path = f"{nusc.dataroot}/{lidar_data['filename']}"
93 | # pcd = LidarPointCloud.from_file(pcd_path).points
94 | #
95 | # # ''' preserve only the inner 64 rings of points (total 128 rings) to reduce pcd size '''
96 | # # mask = pcd[4, :] < 64
97 | # # pcd = pcd[:, mask][:3, :]
98 | # ''' preserve only pts within 40 meter radius '''
99 | # mask = np.sqrt((pcd[0, :]**2 + pcd[1, :]**2)) < 40
100 | # pcd = pcd[:, mask][:3, :]
101 | #
102 | # ego_pose = nusc.get('ego_pose', lidar_data['ego_pose_token'])
103 | # lidar_pose = nusc.get('calibrated_sensor', lidar_data['calibrated_sensor_token'])
104 | # # Get global pose
105 | # ego_pose = transform_matrix(ego_pose['translation'], Quaternion(ego_pose['rotation']), inverse=False)
106 | # lidar_pose = transform_matrix(lidar_pose['translation'], Quaternion(lidar_pose['rotation']), inverse=False)
107 | # global_pose = np.dot(ego_pose, lidar_pose)
108 | #
109 | # pcd = transform_pcd(pcd, global_pose)
110 | #
111 | # location_pcd_list.append(pcd)
112 | #
113 | # pcd_obj.points = o3d.utility.Vector3dVector(pcd)
114 | #
115 | # if start_flag:
116 | # vis1.add_geometry(pcd_obj)
117 | # start_flag = False
118 | # else:
119 | # vis1.update_geometry(pcd_obj)
120 | #
121 | # vis1.poll_events()
122 | # vis1.update_renderer()
123 | # time.sleep(0.005)
124 | #
125 | # if sample["next"] != "":
126 | # sample = nusc.get('sample', sample["next"])
127 | # else:
128 | # break
129 |
130 | '''======================================================================================='''
131 | # sample_token_dict = {}
132 | # frame_offset_dict = {}
133 | #
134 | # start_flag = True
135 | # max_len = 0
136 | # for i in range(nbr_scene):
137 | # sample_token_dict[str(i)] = get_scene_tokens(nusc, i)
138 | # frame_offset_dict[str(i)] = 0
139 | # nbr_sample = len(sample_token_dict[str(i)])
140 | # if nbr_sample > max_len:
141 | # max_len = nbr_sample
142 | #
143 | # frame = 0
144 | # finish_count = 0
145 | # skip_idx = []
146 | # batch = [0]
147 | #
148 | # last_scene_idx = max(batch)
149 | #
150 | # while True:
151 | # # for _ in tqdm(range(max_len)):
152 | # location_pcd_list = []
153 | # location_color_list = []
154 | #
155 | # if frame % 50 == 0 and frame != 0 and len(batch) < 5:
156 | # last_scene_idx += 1
157 | # if last_scene_idx < nbr_scene:
158 | # batch.append(last_scene_idx)
159 | # frame_offset_dict[str(last_scene_idx)] = frame
160 | #
161 | # for idx in range(len(batch)):
162 | # i = batch[idx]
163 | # ''' if scene_i is finished playing, skip it '''
164 | # if i in skip_idx:
165 | # continue
166 | #
167 | # scene_idx = str(i)
168 | # offset = frame_offset_dict[scene_idx]
169 | # if frame - offset >= len(sample_token_dict[scene_idx]):
170 | # finish_count += 1
171 | # skip_idx.append(i)
172 | #
173 | # last_scene_idx += 1
174 | # if last_scene_idx < nbr_scene:
175 | # ''' substitute the finished scene_idx with a new one '''
176 | # target_idx = batch.index(i)
177 | # batch[target_idx] = last_scene_idx
178 | # frame_offset_dict[str(last_scene_idx)] = frame
179 | # print(f"{frame} scene {i} finished, added scene {last_scene_idx}, offset {frame_offset_dict[str(last_scene_idx)]}, showing {batch}")
180 | # else:
181 | # print(f"{frame} scene {i} finished, this is the last scene")
182 | # continue
183 | #
184 | # sample_token = sample_token_dict[scene_idx][frame - offset]
185 | # sample = nusc.get('sample', sample_token)
186 | #
187 | # sample_data_token = sample['data']['LIDAR_FRONT_CENTER']
188 | #
189 | # lidar_data = nusc.get('sample_data', sample_data_token)
190 | # pcd_path = f"{nusc.dataroot}/{lidar_data['filename']}"
191 | # pcd = LidarPointCloud.from_file(pcd_path).points
192 | #
193 | # ''' preserve only pts within 40 meter radius '''
194 | # mask = np.sqrt((pcd[0, :]**2 + pcd[1, :]**2)) < 40
195 | # pcd = pcd[:, mask][:3, :]
196 | #
197 | # ''' transform pcd from local frame to global frame '''
198 | # ego_pose = nusc.get('ego_pose', lidar_data['ego_pose_token'])
199 | # lidar_pose = nusc.get('calibrated_sensor', lidar_data['calibrated_sensor_token'])
200 | # # Get global pose
201 | # ego_pose = transform_matrix(ego_pose['translation'], Quaternion(ego_pose['rotation']), inverse=False)
202 | # lidar_pose = transform_matrix(lidar_pose['translation'], Quaternion(lidar_pose['rotation']), inverse=False)
203 | # global_pose = np.dot(ego_pose, lidar_pose)
204 | #
205 | # pcd = transform_pcd(pcd, global_pose)
206 | # color = np.array([color_list[idx]])/255
207 | # colors = np.repeat(color, len(pcd), axis=0)
208 | #
209 | # location_pcd_list.append(pcd)
210 | # location_color_list.append(colors)
211 | #
212 | # if finish_count == nbr_scene:
213 | # break
214 | #
215 | # pcd_obj.points = o3d.utility.Vector3dVector(np.concatenate(location_pcd_list, axis=0))
216 | #
217 | # pcd_obj.colors = o3d.utility.Vector3dVector(np.concatenate(location_color_list, axis=0))
218 | # if start_flag:
219 | # vis1.add_geometry(pcd_obj)
220 | # start_flag = False
221 | # else:
222 | # vis1.update_geometry(pcd_obj)
223 | #
224 | # vis1.poll_events()
225 | # vis1.update_renderer()
226 | # time.sleep(0.005)
227 | #
228 | # ''' write into video '''
229 | # o3d_screenshot_1 = vis1.capture_screen_float_buffer()
230 | # o3d_screenshot_1 = (255.0 * np.asarray(o3d_screenshot_1)).astype(np.uint8)
231 | # out0.write(o3d_screenshot_1)
232 | #
233 | # frame += 1
234 |
235 | sample_token_dict = {}
236 | frame_offset_dict = {}
237 |
238 | start_flag = True
239 | max_len = 0
240 | for i in range(nbr_scene):
241 | sample_token_dict[str(i)] = get_scene_tokens(nusc, i)
242 | frame_offset_dict[str(i)] = 0
243 | nbr_sample = len(sample_token_dict[str(i)])
244 | if nbr_sample > max_len:
245 | max_len = nbr_sample
246 |
247 | frame = 0
248 | finish_count = 0
249 | skip_idx = []
250 | batch = [0]
251 |
252 | last_scene_idx = max(batch)
253 |
254 | while True:
255 | # for _ in tqdm(range(max_len)):
256 | location_pcd_list = []
257 | location_color_list = []
258 |
259 | # if frame % 250 == 0 and frame != 0 and len(batch) < 5:
260 | # last_scene_idx += 1
261 | # if last_scene_idx < nbr_scene:
262 | # batch.append(last_scene_idx)
263 | # frame_offset_dict[str(last_scene_idx)] = frame
264 |
265 | for idx in range(len(batch)):
266 | i = batch[idx]
267 | ''' if scene_i is finished playing, skip it '''
268 | if i in skip_idx:
269 | continue
270 |
271 | scene_idx = str(i)
272 | offset = frame_offset_dict[scene_idx]
273 | if frame - offset >= len(sample_token_dict[scene_idx]):
274 | finish_count += 1
275 | skip_idx.append(i)
276 |
277 | last_scene_idx += 1
278 | # if last_scene_idx < nbr_scene:
279 | if last_scene_idx < 10:
280 | ''' substitute the finished scene_idx with a new one '''
281 | # if len(batch) == 1:
282 | target_idx = batch.index(i)
283 | batch[target_idx] = last_scene_idx
284 | frame_offset_dict[str(last_scene_idx)] = frame
285 | print(f"{frame} scene {i} finished, showing {batch}")
286 | else:
287 | print(f"{frame} scene {i} finished, this is the last scene")
288 | continue
289 |
290 | sample_token = sample_token_dict[scene_idx][frame - offset]
291 | sample = nusc.get('sample', sample_token)
292 |
293 | sample_data_token = sample['data']['LIDAR_FRONT_CENTER']
294 |
295 | lidar_data = nusc.get('sample_data', sample_data_token)
296 | pcd_path = f"{nusc.dataroot}/{lidar_data['filename']}"
297 | pcd = LidarPointCloud.from_file(pcd_path).points
298 |
299 | ''' preserve only pts within 40 meter radius '''
300 | mask = np.sqrt((pcd[0, :] ** 2 + pcd[1, :] ** 2)) < 40
301 | pcd = pcd[:, mask][:3, :]
302 |
303 | ''' transform pcd from local frame to global frame '''
304 | ego_pose = nusc.get('ego_pose', lidar_data['ego_pose_token'])
305 | lidar_pose = nusc.get('calibrated_sensor', lidar_data['calibrated_sensor_token'])
306 | # Get global pose
307 | ego_pose = transform_matrix(ego_pose['translation'], Quaternion(ego_pose['rotation']), inverse=False)
308 | lidar_pose = transform_matrix(lidar_pose['translation'], Quaternion(lidar_pose['rotation']), inverse=False)
309 | global_pose = np.dot(ego_pose, lidar_pose)
310 |
311 | pcd = transform_pcd(pcd, global_pose)
312 | # color = np.array([color_list[idx]]) / 255
313 | color = np.array([cmap(1/10 * i)[:3]])
314 | # print(f"i {i}, color {1/nbr_scene * i}/1, {color}")
315 | colors = np.repeat(color, len(pcd), axis=0)
316 |
317 | location_pcd_list.append(pcd)
318 | location_color_list.append(colors)
319 |
320 | # if finish_count == nbr_scene:
321 | # break
322 | if finish_count == 10:
323 | break
324 |
325 | if len(location_pcd_list) == 0:
326 | continue
327 |
328 | ''' new lidar scan '''
329 | pcd_obj.points = o3d.utility.Vector3dVector(np.concatenate(location_pcd_list, axis=0))
330 | pcd_obj.colors = o3d.utility.Vector3dVector(np.concatenate(location_color_list, axis=0))
331 |
332 | ''' new background (down-sampled) '''
333 | if start_flag:
334 | background = pcd_obj.voxel_down_sample(voxel_size=4)
335 | background.points = o3d.utility.Vector3dVector(background.points - np.array([0, 0, 2]))
336 | else:
337 | down_sampled = pcd_obj.voxel_down_sample(voxel_size=4)
338 | background.points = o3d.utility.Vector3dVector(np.concatenate((background.points, down_sampled.points - np.array([0, 0, 2])), axis=0))
339 | background.colors = o3d.utility.Vector3dVector(np.concatenate((background.colors, down_sampled.colors), axis=0))
340 | background = background.voxel_down_sample(voxel_size=1)
341 |
342 | ''' new full frame (background + new scan) '''
343 | pcd_obj.points = o3d.utility.Vector3dVector(np.concatenate((background.points, pcd_obj.points), axis=0))
344 | pcd_obj.colors = o3d.utility.Vector3dVector(np.concatenate((background.colors, pcd_obj.colors), axis=0))
345 |
346 | if start_flag:
347 | vis1.add_geometry(pcd_obj)
348 | start_flag = False
349 | else:
350 | vis1.update_geometry(pcd_obj)
351 |
352 | vis1.poll_events()
353 | vis1.update_renderer()
354 | time.sleep(0.005)
355 |
356 | ''' write into video '''
357 | o3d_screenshot_1 = vis1.capture_screen_float_buffer()
358 | o3d_screenshot_1 = (255.0 * np.asarray(o3d_screenshot_1)).astype(np.uint8)
359 | out0.write(o3d_screenshot_1)
360 |
361 | frame += 1
362 |
--------------------------------------------------------------------------------