The train/test pickle files in V2XSim 1.0
CatOneTwo opened this issue · 1 comments
CatOneTwo commented
Hi, thanks for your project.
I notice CoAlign supports the V2XSim2.0, but I don't have enough space to store V2X-Sim 2.0.
Could you please provide the train/test pickle files in V2XSim1.0?
Thank you!
yifanlu0227 commented
Here is the code generating pkl file for V2X-Sim 2.0, which I forget to add to the repo. I think it also works for V2XSim1.0. Further, if you want to use the V2XSim 2.0, you can only download the lidar files and labels. Camera files are not used in this repo.
from typing import OrderedDict
from nuscenes.nuscenes import NuScenes
from icecream import ic
import numpy as np
from pyquaternion import Quaternion
from pathlib import Path
import pickle
import os
from tqdm import tqdm
from opencood.utils.transformation_utils import x1_to_x2, x_to_world, tfm_to_pose
from opencood.utils.box_utils import corner_to_center, mask_boxes_outside_range_numpy, create_bbx, get_points_in_rotated_box_3d
def quaternion_yaw(q: Quaternion) -> float:
"""
Calculate the yaw angle from a quaternion.
Note that this only works for a quaternion that represents a box in lidar or global coordinate frame.
It does not work for a box in the camera frame.
:param q: Quaternion of interest.
:return: Yaw angle in radians.
"""
# Project into xy plane.
v = np.dot(q.rotation_matrix, np.array([1, 0, 0]))
# Measure yaw using arctan.
yaw = np.arctan2(v[1], v[0])
return yaw
def make_split_2(nusc, train_n=80, val_n=10, test_n=10):
""" make split of trainset, valset, testset
for v2x-sim 2.0
"""
scene_num = len(nusc.scene)
assert (train_n + val_n + test_n) == scene_num
# np.random.seed(234)
# perm = np.random.permutation(scene_num)
perm = list(range(100))
train_split = perm[:train_n]
val_split = perm[train_n:train_n+val_n]
test_split = perm[train_n+val_n:train_n+val_n+test_n]
ic(train_split)
ic(val_split)
ic(test_split)
return train_split, val_split, test_split
def make_split_2_new(nusc, train_n=80, val_n=10, test_n=10):
""" make split of trainset, valset, testset
for v2x-sim 2.0
refer to: https://github.com/coperception/coperception/issues/7
"""
scene_num = len(nusc.scene)
assert (train_n + val_n + test_n) == scene_num
# np.random.seed(234)
# perm = np.random.permutation(scene_num)
perm = list(range(100))
train_split = [82,25,95,0,2,6,7,9,10,11,12,13,14,15,16,17,18,20,21,22,23,24,26,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,64,66,67,69,70,71,72,73,74,75,77,80,81,83,85,86,87,88,89,90,93,94,98,99]
val_split = [1,3,4,63,65,68,76,78,79,84]
test_split = [5,8,19,27,28,29,91,92,96,97]
ic(train_split)
ic(val_split)
ic(test_split)
return train_split, val_split, test_split
def build_hash_map(nusc):
instance_set = set()
for sample in nusc.sample:
anns = sample['anns']
instance_token = [nusc.get("sample_annotation",anno_token)['instance_token'] for anno_token in anns]
instance_set.update(instance_token)
hash_map = dict()
for idx, token in enumerate(instance_set):
hash_map[token] = idx
return hash_map
def fill_split_info(nusc, scene_ids, save_filename=None):
""" use the v2x-sim dataset in nuscenes format, and aggregate the point cloud file path,
ego pose, gt boxes information. For one split(train/val/test).
only support single sweep now.
Args:
nusc: nuscenes dataset object
scene_ids: list
"""
split_infos = []
if save_filename:
print(save_filename)
for scene_id in tqdm(scene_ids):
scene = nusc.scene[scene_id]
sample_token = scene['first_sample_token']
while(sample_token != ''):
sample = nusc.get('sample', sample_token) # dict
agent_num = eval(max([i[-1] for i in sample['data'].keys() if i.startswith("LIDAR_TOP")]))
# LIDAR_TOP_id_0 is not vehicle. It's roadside unit.
info = {
'token': sample['token'],
'timestamp': sample['timestamp'],
'agent_num': agent_num,
}
"""
Updated by Yifan Lu, 2022.9.21
1. no shuffle agent's order.
2. filter boxes without lidar point hit.
3. every agent will find it's gt boxes.
"""
# shuffle_order = np.random.permutation(agent_num)
shuffle_order = list(range(1, agent_num+1))
for i in range(agent_num):
idx_in_nuscenes = shuffle_order[i]
idx_in_info = i+1
lidar_sample_data = nusc.get("sample_data", sample['data'][f'LIDAR_TOP_id_{idx_in_nuscenes}']) # dict
info[f'lidar_path_{idx_in_info}'] = nusc.get_sample_data_path(lidar_sample_data['token'])
# dict, include token, timestamp, rotation, translation. rotation is Quaternion
ego_pose_record = nusc.get("ego_pose", lidar_sample_data['ego_pose_token'])
q = Quaternion(ego_pose_record['rotation'])
T_world_ego = q.transformation_matrix
T_world_ego[:3,3] = ego_pose_record['translation']
# info[f'ego_pose_{shuffle_idx}'] = T_world_ego
# dict, include token, timestamp, rotation, translation. rotation is Quaternion
cs_record = nusc.get("calibrated_sensor", lidar_sample_data['calibrated_sensor_token'])
translation_ego_lidar = cs_record['translation']
rotation_ego_lidar = cs_record['rotation']
q = Quaternion(rotation_ego_lidar)
T_ego_lidar = q.transformation_matrix
T_ego_lidar[:3,3] = translation_ego_lidar
T_world_lidar = np.dot(T_world_ego, T_ego_lidar)
info[f'lidar_pose_{idx_in_info}'] = T_world_lidar
# id 0 exists for sure, it will fetch all annotations in the sample
# they are shared for all agents
# boxes are in global coordinate
boxes = nusc.get_boxes(sample['data']['LIDAR_TOP_id_1'])
locs = np.array([b.center for b in boxes]).reshape(-1, 3)
dims = np.array([b.wlh for b in boxes]).reshape(-1, 3)[:, [1, 0, 2]] # wlh == > dxdydz (lwh)
rots = np.array([b.orientation.elements for b in boxes]).reshape(-1, 4)
names = np.array([b.name for b in boxes]) # like "vehicle.audi.tt"
tokens = [b.token for b in boxes]
object_ids = [hash_map[nusc.get("sample_annotation", anno_token)['instance_token']] for anno_token in tokens]
tokens = np.array(tokens)
object_ids = np.array(object_ids)
gt_boxes = np.concatenate([locs, dims, rots], axis=1) # [N, 10]
# filter1: vehicle type
# filter2: size
vehicle_mask1 = [name.startswith("vehicle") for name in names]
vehicle_mask2 = (dims[:,1] > 1.5).tolist() # filter small vehicle, width smaller than 1.5m. Are they really vehicles?
vehicle_mask = np.array([i and j for (i, j) in zip(vehicle_mask1,vehicle_mask2)], dtype=bool)
names = names[vehicle_mask]
tokens = tokens[vehicle_mask]
gt_boxes = gt_boxes[vehicle_mask]
object_ids = object_ids[vehicle_mask]
# filter3: box in range
# filter4: box with lidar point hit.
# load the corresponding data
nbr_dims = 4 # x,y,z,intensity
scan = np.fromfile(info[f'lidar_path_{idx_in_info}'], dtype='float32')
lidar_np = scan.reshape((-1, 5))[:, :nbr_dims] # [N, 4], in ego coord.
lidar_range = [-90, -90, -3, 90, 90, 2]
box_mask = []
for (i, object_content) in enumerate(gt_boxes):
x,y,z,dx,dy,dz,w,a,b,c = object_content
q = Quaternion([w,a,b,c])
T_world_object = q.transformation_matrix
T_world_object[:3,3] = object_content[:3]
object2lidar = np.linalg.solve(T_world_lidar, T_world_object) # T_lidar_object
# shape (3, 8).
# or we can use the create_bbx funcion.
x_corners = dx / 2 * np.array([ 1, 1, -1, -1, 1, 1, -1, -1]) # (8,)
y_corners = dy / 2 * np.array([-1, 1, 1, -1, -1, 1, 1, -1])
z_corners = dz / 2 * np.array([-1, -1, -1, -1, 1, 1, 1, 1])
bbx = np.vstack((x_corners, y_corners, z_corners)) # (3, 8)
# bounding box under ego coordinate shape (4, 8)
bbx = np.r_[bbx, [np.ones(bbx.shape[1])]]
# project the 8 corners to world coordinate
bbx_lidar = np.dot(object2lidar, bbx).T # (8, 4)
bbx_lidar = np.expand_dims(bbx_lidar[:, :3], 0) # (1, 8, 3)
bbx_lidar = corner_to_center(bbx_lidar, order='hwl')
bbx_lidar = mask_boxes_outside_range_numpy(bbx_lidar,
lidar_range,
'hwl')
if bbx_lidar.shape[0] == 0:
box_mask.append(0)
continue
enlarge = 0.2 # 0.2 m
enlarge_extent = [dx/2 + enlarge, dy/2 + enlarge, dz/2 + enlarge]
enlarge_bbx = create_bbx(enlarge_extent).T
# bounding box under ego coordinate shape (4, 8)
enlarge_bbx = np.r_[enlarge_bbx, [np.ones(enlarge_bbx.shape[1])]]
enlarge_bbx_lidar = np.dot(object2lidar, enlarge_bbx).T # (8, 4)
enlarge_bbx_lidar = enlarge_bbx_lidar[:, :3] # (8, 3)
points_in_bbx = get_points_in_rotated_box_3d(lidar_np[:,:3], enlarge_bbx_lidar)
if points_in_bbx.shape[0] == 0:
box_mask.append(0)
else:
box_mask.append(1)
box_mask = np.array(box_mask, dtype=np.bool)
info[f'labels_{idx_in_info}'] = OrderedDict()
info[f'labels_{idx_in_info}']['gt_names'] = names[box_mask]
info[f'labels_{idx_in_info}']['gt_boxes_token'] = tokens[box_mask]
info[f'labels_{idx_in_info}']['gt_boxes_global'] = gt_boxes[box_mask]
info[f'labels_{idx_in_info}']['gt_object_ids'] = object_ids[box_mask]
info["sample_token"] = sample_token
# also store the frame token
split_infos.append(info)
sample_token = sample['next']
if save_filename is not None:
with open(save_filename, 'wb') as f:
pickle.dump(split_infos, f)
else:
return split_infos
def fill_infos(nusc, train_split, val_split, test_split):
""" use the v2x-sim dataset in nuscenes format, and aggregate the point cloud file path,
ego pose, gt boxes information.
Args:
data_path: str, root path for nusc dataset
nusc: nuscenes dataset object
*_split: list, contains index of scene
"""
scene_num = len(nusc.scene)
train_infos = fill_split_info(nusc, train_split)
val_infos = fill_split_info(nusc, val_split)
test_infos = fill_split_info(nusc, test_split)
return train_infos, val_infos, test_infos
def create_pkl(nusc, save_dir):
""" create info.pkl for train/val/test set
Args:
nusc: nuscenes dataset object
save_dir: directory for saving pkl.
Returns: None
"""
train_split, val_split, test_split = make_split_2_new(nusc)
if MULTI_PROCESSES:
fill_infos_multiprocess(nusc, train_split, val_split, test_split, save_dir)
return
train_infos, val_infos, test_infos = fill_infos(nusc, train_split, val_split, test_split)
print(f'train sample: {len(train_infos)}, val sample: {len(val_infos)}, test sample: {len(test_infos)}')
with open(os.path.join(save_dir, f'v2xsim_infos_train.pkl'), 'wb') as f:
pickle.dump(train_infos, f)
with open(os.path.join(save_dir, f'v2xsim_infos_val.pkl'), 'wb') as f:
pickle.dump(val_infos, f)
with open(os.path.join(save_dir, f'v2xsim_infos_test.pkl'), 'wb') as f:
pickle.dump(test_infos, f)
if __name__=="__main__":
nusc = NuScenes(version='v1.0-mini', dataroot='/GPFS/rhome/yifanlu/workspace/dataset/v2xsim2-complete', verbose=True) # specify your own path.
output_path = "/GPFS/rhome/yifanlu/workspace/dataset/v2xsim2_info_no_shuffle_agent" # a folder path. specify your own path.
global hash_map
hash_map = build_hash_map(nusc)
create_pkl(nusc, output_path)