def __init__(self, config={}): self.config = { 'version': 'v1.0-mini', 'dataroot': os.path.join(os.environ['PKG_PATH'],'data') } self.config.update(config) self.nusc = NuScenes(version=self.config['version'], dataroot=self.config['dataroot'], verbose=True) self.nusc_can = NuScenesCanBus(dataroot=self.config['dataroot']) self.utils = Utils(self.nusc, self.nusc_can)
def init_nu(): """ ------------------------------------------------------------------------------------------------------------- Initialize nuScene return: [tuple] NuScenes instance, NuScenesCanBus instance ------------------------------------------------------------------------------------------------------------- """ nu = NuScenes( version=nuscenes_ver, dataroot=dir_orig, verbose=verbose ) # NuScenes database nc = NuScenesCanBus( dataroot=dir_orig ) # NuScenes CAN bus return nu, nc
def process_once(scene_name_list=[], data_save_dir=None, config={}): """get raw data from dataset, everything are lists :param scene_name_list: list of scene names to process :param data_save_dir: directory to save processed raw data :param config: additional configs :returns: one pandas dataframe for each scene """ test_scene_name = config['scene_name'] if config['num_workers'] == 1: nusc = config['other_configs']['nusc'] helper = config['other_configs']['helper'] rasterizer = config['other_configs']['rasterizer'] else: nusc = ray.get(config['other_configs']['nusc']) helper = ray.get(config['other_configs']['helper']) rasterizer = ray.get(config['other_configs']['rasterizer']) dataroot = config['other_configs']['dataroot'] nusc_can = NuScenesCanBus(dataroot=dataroot) df_dict = { #### scene info #### 'scene_name': [], 'scene_token':[], 'scene_nbr_samples':[], 'scene_description':[], 'scene_log_token':[], 'scene_map_token':[], 'scene_vehicle_id': [], 'scene_date_captured':[], 'scene_location':[], #### sample info #### 'sample_idx': [], 'sample_token':[], 'time_stamp': [], #### ego info #### 'current_ego_pos': [], 'current_ego_quat':[], 'current_ego_accel':[], 'current_ego_speed':[], 'current_ego_steering': [], 'current_ego_on_road_objects':[], 'current_ego_raster_img':[], 'current_ego_raster_img_path': [], #### instance info #### 'annotation_token': [], 'instance_token': [], 'instance_category': [], 'instance_attributes': [], 'instance_nbr_annotations':[], 'current_instance_pos': [], 'current_instance_quat': [], 'current_instance_speed': [], 'current_instance_accel':[], 'current_instance_on_road_objects':[], 'current_instance_raster_img':[], 'current_instance_raster_img_path': [] } #### loop over scenes #### for scene_name in tqdm.tqdm(scene_name_list, 'processing scene'): if test_scene_name is not None: if not scene_name == test_scene_name: continue scene_df_dict = copy.deepcopy(df_dict) scene_token = nusc.field2token('scene', 'name', scene_name)[0] scene = nusc.get('scene', scene_token) scene_log = nusc.get('log', scene['log_token']) scene_map = nusc.get('map', scene_log['map_token']) nusc_map = NuScenesMap(dataroot=dataroot, map_name=scene_log['location']) raster_dir = os.path.join(data_save_dir, scene_name) if not os.path.isdir(raster_dir): os.makedirs(raster_dir) # can bus information if scene_name in NO_CANBUS_SCENES: ego_accel_traj = [None] * 50 ego_rotation_rate_traj = [None] * 50 ego_speed_traj = [None] * 50 else: pose_msg = nusc_can.get_messages(scene_name, 'pose') ego_accel_traj = [pm['accel'] for pm in pose_msg][::25] ego_rotation_rate_traj = [pm['rotation_rate'] for pm in pose_msg][::25] ego_speed_traj = [pm['vel'] for pm in pose_msg][::25] #### loop over samples #### sample_token = scene['first_sample_token'] sample = nusc.get('sample', sample_token) sample_idx = 0 t = time.time() while sample['next'] != '': # print(sample_idx) # print(t-time.time()) # t = time.time() cam_data = nusc.get('sample_data', sample['data']['CAM_FRONT']) ego_pose = nusc.get('ego_pose', cam_data['ego_pose_token']) # ego raster ego_raster = rasterizer.make_input_representation(instance_token=None, sample_token=sample_token, ego=True, ego_pose=ego_pose, include_history=False) ego_rel_path = str(sample_idx) + "_" + 'ego.png' current_ego_raster_img_path = os.path.join(raster_dir, ego_rel_path) ego_raster_png = im.fromarray(ego_raster) ego_raster_png.save(current_ego_raster_img_path) # ego road objects ego_on_road_objects = nusc_map.layers_on_point(ego_pose['translation'][0], ego_pose['translation'][1]) #### loop over annotations #### for ann_token in sample['anns']: ann = nusc.get('sample_annotation', ann_token) instance = nusc.get('instance', ann['instance_token']) if filter_instance(ann, nusc): #### populate data #### ## scene ## scene_df_dict['scene_name'].append(scene_name) scene_df_dict['scene_token'].append(scene_token) scene_df_dict['scene_nbr_samples'].append(scene['nbr_samples']) scene_df_dict['scene_description'].append(scene['description']) scene_df_dict['scene_log_token'].append(scene_log['token']) scene_df_dict['scene_map_token'].append(scene_map['token']) scene_df_dict['scene_vehicle_id'].append(scene_log['vehicle']) scene_df_dict['scene_date_captured'].append(scene_log['date_captured']) scene_df_dict['scene_location'].append(scene_log['location']) ## sample info ## scene_df_dict['sample_idx'].append(sample_idx) scene_df_dict['sample_token'].append(sample_token) scene_df_dict['time_stamp'].append(sample['timestamp']) ## ego info ## scene_df_dict['current_ego_pos'].append(ego_pose['translation']) scene_df_dict['current_ego_quat'].append(ego_pose['rotation']) idx = min(sample_idx, len(ego_speed_traj)-1) scene_df_dict['current_ego_speed'].append(ego_speed_traj[idx]) idx = min(sample_idx, len(ego_accel_traj)-1) scene_df_dict['current_ego_accel'].append(ego_accel_traj[idx]) idx = min(sample_idx, len(ego_rotation_rate_traj)-1) scene_df_dict['current_ego_steering'].append(ego_rotation_rate_traj[idx]) scene_df_dict['current_ego_on_road_objects'].append(ego_on_road_objects) scene_df_dict['current_ego_raster_img'].append(ego_raster) scene_df_dict['current_ego_raster_img_path'].append(scene_name+"/"+ego_rel_path) ## instance info ## scene_df_dict['annotation_token'].append(ann_token) scene_df_dict['instance_token'].append(ann['instance_token']) scene_df_dict['instance_category'].append(ann['category_name']) instance_attributes = [nusc.get('attribute', attribute_token)['name'] for attribute_token in ann['attribute_tokens']] scene_df_dict['instance_attributes'].append(instance_attributes) scene_df_dict['instance_nbr_annotations'].append(instance['nbr_annotations']) scene_df_dict['current_instance_pos'].append(ann['translation']) scene_df_dict['current_instance_quat'].append(ann['rotation']) instance_speed = helper.get_velocity_for_agent(ann['instance_token'], sample_token) scene_df_dict['current_instance_speed'].append(instance_speed) instance_accel = helper.get_acceleration_for_agent(ann['instance_token'], sample_token) scene_df_dict['current_instance_accel'].append(instance_accel) instance_on_road_objects = nusc_map.layers_on_point(ann['translation'][0], ann['translation'][1]) scene_df_dict['current_instance_on_road_objects'].append(instance_on_road_objects) instance_raster = rasterizer.make_input_representation(instance_token=ann['instance_token'], sample_token=sample_token, ego=False, include_history=False) scene_df_dict['current_instance_raster_img'].append(instance_raster) instance_rel_path = str(sample_idx) + "_"+ann['instance_token']+'.png' instance_raster_img_path = os.path.join(raster_dir, instance_rel_path) instance_raster_png = im.fromarray(instance_raster) instance_raster_png.save(instance_raster_img_path) scene_df_dict['current_instance_raster_img_path'].append(scene_name+"/"+instance_rel_path) sample_token = sample['next'] sample = nusc.get('sample', sample_token) sample_idx += 1 scene_df = pd.DataFrame(scene_df_dict) scene_df.to_pickle(data_save_dir+"/"+scene_name+".pkl")
def __init__(self, nusc, nusc_split, kwargs, seed=0): super(nuScenesDataset, self).__init__() # set seed for split np.random.seed(seed) self.nusc = nusc self.nusc_root = self.nusc.dataroot self.nusc_can = NuScenesCanBus(dataroot=self.nusc_root) self.nusc_split = nusc_split # number of input samples self.n_input = kwargs["n_input"] # number of sampled trajectories self.n_samples = kwargs["n_samples"] # number of output samples self.n_output = kwargs["n_output"] assert(self.n_output == 7) # self.train_on_all_sweeps = kwargs["train_on_all_sweeps"] # scene-0419 does not have vehicle monitor data blacklist = [419] + self.nusc_can.can_blacklist # NOTE: use the official split (minus the ones in the blacklist) if "scene_token" in kwargs and kwargs["scene_token"] != "": scene = self.nusc.get("scene", kwargs["scene_token"]) scenes = [scene] else: scene_splits = create_splits_scenes(verbose=False) scene_names = scene_splits[self.nusc_split] scenes = [] for scene in self.nusc.scene: scene_name = scene["name"] scene_no = int(scene_name[-4:]) if (scene_name in scene_names) and (scene_no not in blacklist): scenes.append(scene) # list all sample data self.valid_index = [] self.flip_flags = [] self.scene_tokens = [] self.sample_data_tokens = [] for scene in scenes: scene_token = scene["token"] # location log = self.nusc.get("log", scene["log_token"]) # flip x axis if in left-hand traffic (singapore) flip_flag = True if log["location"].startswith("singapore") else False # record the token of every key frame start_index = len(self.sample_data_tokens) first_sample = self.nusc.get("sample", scene["first_sample_token"]) sample_data_token = first_sample["data"]["LIDAR_TOP"] while sample_data_token != "": sample_data = self.nusc.get("sample_data", sample_data_token) if (self.nusc_split == "train" and self.train_on_all_sweeps) or (sample_data["is_key_frame"]): self.flip_flags.append(flip_flag) self.scene_tokens.append(scene_token) self.sample_data_tokens.append(sample_data_token) sample_data_token = sample_data["next"] end_index = len(self.sample_data_tokens) # NOTE: make sure we have enough number of sweeps for input and output if self.nusc_split == "train" and self.train_on_all_sweeps: valid_start_index = start_index + self.n_input - 1 valid_end_index = end_index - (self.n_output - 1) * self.N_SWEEPS_PER_SAMPLE else: # NEW: acknowledge the fact and skip the first sample n_input_samples = self.n_input // self.N_SWEEPS_PER_SAMPLE valid_start_index = start_index + n_input_samples valid_end_index = end_index - self.n_output + 1 self.valid_index += list(range(valid_start_index, valid_end_index)) self._n_examples = len(self.valid_index) print(f"{self.nusc_split}: {self._n_examples} valid samples over {len(scenes)} scenes")
class nuScenesDataset(Dataset): N_SWEEPS_PER_SAMPLE = 10 SAMPLE_INTERVAL = 0.5 # second def __init__(self, nusc, nusc_split, kwargs, seed=0): super(nuScenesDataset, self).__init__() # set seed for split np.random.seed(seed) self.nusc = nusc self.nusc_root = self.nusc.dataroot self.nusc_can = NuScenesCanBus(dataroot=self.nusc_root) self.nusc_split = nusc_split # number of input samples self.n_input = kwargs["n_input"] # number of sampled trajectories self.n_samples = kwargs["n_samples"] # number of output samples self.n_output = kwargs["n_output"] assert(self.n_output == 7) # self.train_on_all_sweeps = kwargs["train_on_all_sweeps"] # scene-0419 does not have vehicle monitor data blacklist = [419] + self.nusc_can.can_blacklist # NOTE: use the official split (minus the ones in the blacklist) if "scene_token" in kwargs and kwargs["scene_token"] != "": scene = self.nusc.get("scene", kwargs["scene_token"]) scenes = [scene] else: scene_splits = create_splits_scenes(verbose=False) scene_names = scene_splits[self.nusc_split] scenes = [] for scene in self.nusc.scene: scene_name = scene["name"] scene_no = int(scene_name[-4:]) if (scene_name in scene_names) and (scene_no not in blacklist): scenes.append(scene) # list all sample data self.valid_index = [] self.flip_flags = [] self.scene_tokens = [] self.sample_data_tokens = [] for scene in scenes: scene_token = scene["token"] # location log = self.nusc.get("log", scene["log_token"]) # flip x axis if in left-hand traffic (singapore) flip_flag = True if log["location"].startswith("singapore") else False # record the token of every key frame start_index = len(self.sample_data_tokens) first_sample = self.nusc.get("sample", scene["first_sample_token"]) sample_data_token = first_sample["data"]["LIDAR_TOP"] while sample_data_token != "": sample_data = self.nusc.get("sample_data", sample_data_token) if (self.nusc_split == "train" and self.train_on_all_sweeps) or (sample_data["is_key_frame"]): self.flip_flags.append(flip_flag) self.scene_tokens.append(scene_token) self.sample_data_tokens.append(sample_data_token) sample_data_token = sample_data["next"] end_index = len(self.sample_data_tokens) # NOTE: make sure we have enough number of sweeps for input and output if self.nusc_split == "train" and self.train_on_all_sweeps: valid_start_index = start_index + self.n_input - 1 valid_end_index = end_index - (self.n_output - 1) * self.N_SWEEPS_PER_SAMPLE else: # NEW: acknowledge the fact and skip the first sample n_input_samples = self.n_input // self.N_SWEEPS_PER_SAMPLE valid_start_index = start_index + n_input_samples valid_end_index = end_index - self.n_output + 1 self.valid_index += list(range(valid_start_index, valid_end_index)) self._n_examples = len(self.valid_index) print(f"{self.nusc_split}: {self._n_examples} valid samples over {len(scenes)} scenes") def __len__(self): return self._n_examples def get_global_pose(self, sd_token, inverse=False): sd = self.nusc.get("sample_data", sd_token) sd_ep = self.nusc.get("ego_pose", sd["ego_pose_token"]) sd_cs = self.nusc.get("calibrated_sensor", sd["calibrated_sensor_token"]) if inverse is False: global_from_ego = transform_matrix(sd_ep["translation"], Quaternion(sd_ep["rotation"]), inverse=False) ego_from_sensor = transform_matrix(sd_cs["translation"], Quaternion(sd_cs["rotation"]), inverse=False) pose = global_from_ego.dot(ego_from_sensor) else: sensor_from_ego = transform_matrix(sd_cs["translation"], Quaternion(sd_cs["rotation"]), inverse=True) ego_from_global = transform_matrix(sd_ep["translation"], Quaternion(sd_ep["rotation"]), inverse=True) pose = sensor_from_ego.dot(ego_from_global) return pose def load_ground_segmentation(self, sample_data_token): path = f"{self.nusc.dataroot}/grndseg/{self.nusc.version}/{sample_data_token}_grndseg.bin" labels = np.fromfile(path, np.uint8) return labels def load_future_visible_freespace(self, sample_data_token): path = f"{self.nusc.dataroot}/fvfmaps/{self.nusc.version}/{sample_data_token}.bin" if os.path.exists(path): fvf_maps = np.fromfile(path, dtype=np.int8) fvf_maps = fvf_maps.reshape((7, 704, 400)) else: fvf_maps = np.zeros((7, 704, 400), dtype=np.int8) warnings.warn(f"Cannot find fvf_maps at {path}") return fvf_maps def load_object_boxes(self, sample_data_token): path = f"{self.nusc.dataroot}/obj_boxes/{self.nusc.version}/{sample_data_token}.bin" if os.path.exists(path): obj_boxes = np.fromfile(path, dtype=bool) obj_boxes = obj_boxes.reshape((7, 704, 400)) else: obj_boxes = np.zeros((7, 704, 400)) return obj_boxes def load_object_shadows(self, sample_data_token): path = f"{self.nusc.dataroot}/obj_shadows/{self.nusc.version}/{sample_data_token}.bin" if os.path.exists(path): obj_shadows = np.fromfile(path, dtype=bool) obj_shadows = obj_shadows.reshape((7, 704, 400)) else: obj_shadows = np.zeros((7, 704, 400)) return obj_shadows def __getitem__(self, idx): ref_index = self.valid_index[idx] ref_sd_token = self.sample_data_tokens[ref_index] ref_scene_token = self.scene_tokens[ref_index] flip_flag = self.flip_flags[ref_index] # reference coordinate frame ref_from_global = self.get_global_pose(ref_sd_token, inverse=True) # NOTE: input input_sds = [] sd_token = ref_sd_token while sd_token != "" and len(input_sds) < self.n_input: curr_sd = self.nusc.get("sample_data", sd_token) input_sds.append(curr_sd) sd_token = curr_sd["prev"] # call out when we have less than the desired number of input sweeps # if len(input_sds) < self.n_input: # warnings.warn(f"The number of input sweeps {len(input_sds)} is less than {self.n_input}.", RuntimeWarning) # get input sweep frames input_points_list = [] for i, curr_sd in enumerate(input_sds): # load the current lidar sweep curr_lidar_pc = MyLidarPointCloud.from_file(f"{self.nusc_root}/{curr_sd['filename']}") # remove ego returns curr_lidar_pc.remove_ego() # transform from the current lidar frame to global and then to the reference lidar frame global_from_curr = self.get_global_pose(curr_sd["token"], inverse=False) ref_from_curr = ref_from_global.dot(global_from_curr) curr_lidar_pc.transform(ref_from_curr) # NOTE: check if we are in Singapore (if so flip x) if flip_flag: curr_lidar_pc.points[0] *= -1 # points = np.asarray(curr_lidar_pc.points[:3].T) tindex = np.full((len(points), 1), i) points = np.concatenate((points, tindex), axis=1) # input_points_list.append(points.astype(np.float32)) # NOTE: output # get output sample frames and ground truth trajectory output_origin_list = [] output_points_list = [] gt_trajectory = np.zeros((self.n_output, 3), np.float64) for i in range(self.n_output): if self.nusc_split == "train" and self.train_on_all_sweeps: index = ref_index + i * self.N_SWEEPS_PER_SAMPLE else: index = ref_index + i # if this exists a valid target if index < len(self.scene_tokens) and self.scene_tokens[index] == ref_scene_token: curr_sd_token = self.sample_data_tokens[index] curr_sd = self.nusc.get("sample_data", curr_sd_token) # load the current lidar sweep curr_lidar_pc = LidarPointCloud.from_file(f"{self.nusc_root}/{curr_sd['filename']}") # transform from the current lidar frame to global and then to the reference lidar frame global_from_curr = self.get_global_pose(curr_sd_token, inverse=False) ref_from_curr = ref_from_global.dot(global_from_curr) curr_lidar_pc.transform(ref_from_curr) # theta = quaternion_yaw(Quaternion(matrix=ref_from_curr)) # NOTE: check if we are in Singapore (if so flip x) if flip_flag: ref_from_curr[0, 3] *= -1 curr_lidar_pc.points[0] *= -1 theta *= -1 origin = np.array(ref_from_curr[:3, 3]) points = np.array(curr_lidar_pc.points[:3].T) gt_trajectory[i, :] = [origin[0], origin[1], theta] tindex = np.full(len(points), i) labels = self.load_ground_segmentation(curr_sd_token) assert(len(labels) == len(points)) mask = np.logical_and(labels >= 1, labels <= 30) points = np.concatenate((points, tindex[:, None], labels[:, None]), axis=1) points = points[mask, :] else: # filler raise RuntimeError(f"The {i}-th output frame is not available") origin = np.array([0.0, 0.0, 0.0]) points = np.full((0, 5), -1) # origin output_origin_list.append(origin.astype(np.float32)) # points output_points_list.append(points.astype(np.float32)) # NOTE: trajectory sampling ref_scene = self.nusc.get("scene", ref_scene_token) # NOTE: rely on pose and steeranglefeedback data instead of vehicle_monitor vm_msgs = self.nusc_can.get_messages(ref_scene["name"], "vehicle_monitor") vm_uts = [msg["utime"] for msg in vm_msgs] pose_msgs = self.nusc_can.get_messages(ref_scene["name"], "pose") pose_uts = [msg["utime"] for msg in pose_msgs] steer_msgs = self.nusc_can.get_messages(ref_scene["name"], "steeranglefeedback") steer_uts = [msg["utime"] for msg in steer_msgs] # locate the closest message by universal timestamp ref_sd = self.nusc.get("sample_data", ref_sd_token) ref_utime = ref_sd["timestamp"] vm_index = locate_message(vm_uts, ref_utime) vm_data = vm_msgs[vm_index] pose_index = locate_message(pose_uts, ref_utime) pose_data = pose_msgs[pose_index] steer_index = locate_message(steer_uts, ref_utime) steer_data = steer_msgs[steer_index] # initial speed # v0 = vm_data["vehicle_speed"] / 3.6 # km/h to m/s v0 = pose_data["vel"][0] # [0] means longitudinal velocity # curvature (positive: turn left) # steering = np.deg2rad(vm_data["steering"]) steering = steer_data["value"] if flip_flag: steering *= -1 Kappa = 2 * steering / 2.588 # left_signal = vm_data["left_signal"] right_signal = vm_data["right_signal"] if flip_flag: left_signal, right_signal = right_signal, left_signal drive_command = [left_signal, right_signal] # initial state T0 = np.array([0.0, 1.0]) # define front N0 = np.array([1.0, 0.0]) if Kappa <= 0 else np.array([-1.0, 0.0]) # define side # # tt = np.arange(self.n_output) * self.SAMPLE_INTERVAL # tt = np.arange(0, self.n_output + self.SAMPLE_INTERVAL, self.SAMPLE_INTERVAL) t_start = 0 # second t_end = (self.n_output-1) * self.SAMPLE_INTERVAL # second t_interval = self.SAMPLE_INTERVAL / 10 tt = np.arange(t_start, t_end + t_interval, t_interval) sampled_trajectories_fine = trajectory_sampler.sample(v0, Kappa, T0, N0, tt, self.n_samples) sampled_trajectories = sampled_trajectories_fine[:, ::10] # obj_boxes = self.load_object_boxes(ref_sd_token) obj_shadows = self.load_object_shadows(ref_sd_token) # fvf_maps = self.load_future_visible_freespace(ref_sd_token) # example = { "scene_token": ref_scene_token, "sample_data_token": ref_sd_token, "input_points": torch.from_numpy(np.concatenate(input_points_list)), "sampled_trajectories_fine": torch.from_numpy(sampled_trajectories_fine), "sampled_trajectories": torch.from_numpy(sampled_trajectories), "drive_command": torch.tensor(drive_command), "output_origin": torch.from_numpy(np.stack(output_origin_list)), "output_points": torch.from_numpy(np.concatenate(output_points_list)), "gt_trajectory": torch.from_numpy(gt_trajectory), "obj_boxes": torch.from_numpy(obj_boxes), "obj_shadows": torch.from_numpy(obj_shadows), "fvf_maps": torch.from_numpy(fvf_maps), } return example
def get_can_data(scene_name, data_root): # This py file has a function inside called 'get_can_data(scene_name, data_root)'. # It requires two inputs (both strings like 'blah') of the scene_name (e.g. 'scene_0565') and data_root (e.g. '../data/sets/nuscenes'). The can_bus data (as .json files) must be sitting in a can_bus folder in data/sets/nuscenes. The two dots before data are to go up a level. The can_bus data is kept in this folder in the repo but is ignored when pushing up to GitHub as it is a large sized folder than what GitHub would recommend. # To bring the internal function 'get_can_data' into the workspace, run the following code: # exec(open("../python-functions/can_data_fcn.py").read()) # This will allow you to use get_can_data function as you please. An example of how it is used is located in python-scripts/test_script.py. # When you get the data as a variable name of your choice, the breakdown of the data is as follows: # a. Pick the CAN Message you want. # b. Pick the CAN Signal you want. # ii. veh_msgs: utimes, fl_wheel_speed, fr_wheel_speed, rl_wheel_speed, rr_wheel_speed # pose_msgs: utimes, veh_speed, x, y, rot_quaternion_mat # imu_msgs: utimes, accel_x, accel_y, accel_z, rot_quaternion_mat, roll_rate, pitch_rate, yaw_rate # veh_monitor_msgs: utimes, swa_data, swa_rate_data, veh_speed, yaw_rate # For the data frequency (time-step): # veh_msgs: 100 Hz # pose_msgs: 50 Hz # imu_msgs: 100 Hz # vehicle_monitor_msgs: 2 Hz from nuscenes.can_bus.can_bus_api import NuScenesCanBus import numpy as np import scipy.io as sio import math import matplotlib.pyplot as plt # nusc_can = NuScenesCanBus(dataroot='../../../../../../../Volumes/Arnie SSD/data/sets/nuscenes') nusc_can = NuScenesCanBus(dataroot = data_root) #scene_name = 'scene-0574' ############# VEHICLE INFO MESSAGE ################################## # Obtain time array of vehicle info message veh_messages = nusc_can.get_messages(scene_name, 'zoe_veh_info') veh_utimes = np.array([m['utime'] for m in veh_messages]) veh_msgs_utimes = (veh_utimes - min(veh_utimes)) / 1e6 # Units: sec # Obtain the four different wheel speed sensor data veh_fl_wheel_speed_data = np.array([m['FL_wheel_speed'] for m in veh_messages]) # Units: rpm veh_fr_wheel_speed_data = np.array([m['FR_wheel_speed'] for m in veh_messages]) # Units: rpm veh_rl_wheel_speed_data = np.array([m['RL_wheel_speed'] for m in veh_messages]) # Units: rpm veh_rr_wheel_speed_data = np.array([m['RR_wheel_speed'] for m in veh_messages]) # Units: rpm # Convert wheel speed data to linear velocity for each wheel veh_r_wheel = 0.305; # Units: m veh_msgs_fl_wheel_speed_data = veh_fl_wheel_speed_data * 2 * math.pi * veh_r_wheel / 60 # Units: mps veh_msgs_fr_wheel_speed_data = veh_fr_wheel_speed_data * 2 * math.pi * veh_r_wheel / 60 # Units: mps veh_msgs_rl_wheel_speed_data = veh_rl_wheel_speed_data * 2 * math.pi * veh_r_wheel / 60 # Units: mps veh_msgs_rr_wheel_speed_data = veh_rr_wheel_speed_data * 2 * math.pi * veh_r_wheel / 60 # Units: mps # Obtain steering wheel angle (swa) sensor data veh_msgs_swa_data = np.array([m['steer_corrected'] for m in veh_messages]) # Units: deg veh_msgs = { 'utimes': veh_msgs_utimes, 'fl_wheel_speed': veh_msgs_fl_wheel_speed_data, 'fr_wheel_speed': veh_msgs_fr_wheel_speed_data, 'rl_wheel_speed': veh_msgs_rl_wheel_speed_data, 'rr_wheel_speed': veh_msgs_rr_wheel_speed_data, } ############# POSE MESSAGE ################################## # Obtain time array of pose message pose_messages = nusc_can.get_messages(scene_name, 'pose') pose_utimes = np.array([m['utime'] for m in pose_messages]) pose_msgs_utimes = (pose_utimes - min(pose_utimes)) / 1e6 # Units: sec # Obtain vehicle speed; in the ego vehicle frame pose_msgs_veh_speed = np.array([m['vel'] for m in pose_messages]) # Units: mps # Obtain x position from southeast corner of the map pose_msgs_x = np.array([m['pos'][0] for m in pose_messages]) # Units: m # Obtain y position from southeast corner of the map pose_msgs_y = np.array([m['pos'][1] for m in pose_messages]) # Units: m # Obtain Quaternion Rotation matrix - outputs as an array; in the ego vehicle frame pose_msgs_rot_quaternion_mat = np.array([m['orientation'] for m in pose_messages]) # Units: dimension-less pose_msgs = { 'utimes': pose_msgs_utimes, 'veh_speed': pose_msgs_veh_speed, 'x': pose_msgs_x, 'y': pose_msgs_y, 'rot_quaternion_mat': pose_msgs_rot_quaternion_mat, } ############# IMU MESSAGE ################################## # Obtain time array of IMU message imu_messages = nusc_can.get_messages(scene_name, 'ms_imu') imu_utimes = np.array([m['utime'] for m in imu_messages]) imu_msgs_utimes = (imu_utimes - min(imu_utimes)) / 1e6 # Units: sec # Obtain Acceleration in x direction imu_msgs_accel_x = np.array([m['linear_accel'][0] for m in imu_messages]) # Units: m/s/s # Obtain Acceleration in y direction imu_msgs_accel_y = np.array([m['linear_accel'][1] for m in imu_messages]) # Units: m/s/s # Obtain Acceleration in z direction imu_msgs_accel_z = np.array([m['linear_accel'][2] for m in imu_messages]) # Units: m/s/s # Obtain Quaternion Rotation matrix - outputs as an array imu_msgs_rot_quaternion_mat = np.array([m['q'] for m in imu_messages]) # Units: dimension-less # Obtain Roll Rate imu_msgs_roll_rate = np.array([m['rotation_rate'][0] for m in imu_messages]) # Units: rad/s # Obtain Pitch Rate imu_msgs_pitch_rate = np.array([m['rotation_rate'][1] for m in imu_messages]) # Units: rad/s # Obtain Yaw Rate imu_msgs_yaw_rate = np.array([m['rotation_rate'][2] for m in imu_messages]) # Units: rad/s imu_msgs = { 'utimes': imu_msgs_utimes, 'accel_x': imu_msgs_accel_x, 'accel_y': imu_msgs_accel_y, 'accel_z': imu_msgs_accel_z, 'rot_quaternion_mat': imu_msgs_rot_quaternion_mat, 'roll_rate': imu_msgs_roll_rate, 'pitch_rate': imu_msgs_pitch_rate, 'yaw_rate': imu_msgs_yaw_rate, } ############# VEHICLE MONITOR MESSAGE ################################## # Note this is different than the Vehicle Info message. This is recorded at a slower frequency. # Obtain time array of Vehicle Monitor message veh_monitor_messages = nusc_can.get_messages(scene_name, 'vehicle_monitor') veh_monitor_utimes = np.array([m['utime'] for m in veh_monitor_messages]) veh_monitor_msgs_utimes = (veh_monitor_utimes - min(veh_monitor_utimes)) / 1e6 # Units: sec # Obtain steering wheel angle (swa) sensor data veh_monitor_msgs_swa_data = np.array([m['steering'] for m in veh_monitor_messages]) # Units: deg # Obtain steering wheel angle rate (swa_dot) sensor data veh_monitor_msgs_swa_rate_data = np.array([m['steering_speed'] for m in veh_monitor_messages]) # Units: deg/s # Obtain vehicle speed sensor data veh_monitor_msgs_veh_speed = np.array([m['vehicle_speed'] for m in veh_monitor_messages]) # Units: km/h # Obtain yaw rate sensor data veh_monitor_msgs_yaw_rate = np.array([m['yaw_rate'] for m in veh_monitor_messages]) # Units: deg/s veh_monitor_msgs = { 'utimes': veh_monitor_msgs_utimes, 'swa_data': veh_monitor_msgs_swa_data, 'swa_rate_data': veh_monitor_msgs_swa_rate_data, 'veh_speed': veh_monitor_msgs_veh_speed, 'yaw_rate': veh_monitor_msgs_yaw_rate, } # Pass all structures into can_data_out can_data_out = { 'veh_msgs': veh_msgs, 'pose_msgs': pose_msgs, 'imu_msgs': imu_msgs, 'veh_monitor_msgs': veh_monitor_msgs, } # Save the long_accel data into a mat file #sio.savemat('np_vector.mat', {'long_accel':data}) # Plot mps wheel speeds and mps vehicle speed together #plt.plot(pose_x, pose_y) #plt.show() #nusc_can.plot_baseline_route(scene_name) return can_data_out
def create_nuscenes_infos(root_path, modality, radar_version, version="v1.0-trainval", max_sweeps=10): from nuscenes.nuscenes import NuScenes from nuscenes.can_bus.can_bus_api import NuScenesCanBus nusc = NuScenes(version=version, dataroot=root_path, verbose=True) nusc_can = NuScenesCanBus(dataroot=root_path) from nuscenes.utils import splits available_vers = ["v1.0-trainval", "v1.0-test", "v1.0-mini"] assert version in available_vers if version == "v1.0-trainval": train_scenes = splits.train val_scenes = splits.val elif version == "v1.0-test": train_scenes = splits.test val_scenes = [] elif version == "v1.0-mini": train_scenes = splits.mini_train val_scenes = splits.mini_val else: raise ValueError("unknown") test = "test" in version root_path = Path(root_path) output_path = Path('/home/spalab/jskim/second.pytorch/second/dataset') # filter exist scenes. you may only download part of dataset. available_scenes = _get_available_scenes(nusc, nusc_can) available_scene_names = [s["name"] for s in available_scenes] train_scenes = list( filter(lambda x: x in available_scene_names, train_scenes)) val_scenes = list(filter(lambda x: x in available_scene_names, val_scenes)) train_scenes = set([ available_scenes[available_scene_names.index(s)]["token"] for s in train_scenes ]) val_scenes = set([ available_scenes[available_scene_names.index(s)]["token"] for s in val_scenes ]) if test: print(f"test scene: {len(train_scenes)}") else: print( f"train scene: {len(train_scenes)}, val scene: {len(val_scenes)}") train_nusc_infos, val_nusc_infos = _fill_trainval_infos( nusc, nusc_can, train_scenes, val_scenes, modality, radar_version, test=test, max_sweeps=max_sweeps) metadata = { "version": version, } if test: print(f"test sample: {len(train_nusc_infos)}") data = { "infos": train_nusc_infos, "metadata": metadata, } with open(output_path / "infos_test.pkl", 'wb') as f: pickle.dump(data, f) else: print( f"train sample: {len(train_nusc_infos)}, val sample: {len(val_nusc_infos)}" ) data = { "infos": train_nusc_infos, "metadata": metadata, } if max_sweeps == 6: with open( output_path / f"infos_train_multisweep_{radar_version}.pkl", 'wb') as f: pickle.dump(data, f) data["infos"] = val_nusc_infos with open( output_path / f"infos_val_multisweep_{radar_version}.pkl", 'wb') as f: pickle.dump(data, f) elif max_sweeps == 10: with open( output_path / f"infos_train_multisweep_{radar_version}_sweeps{max_sweeps}.pkl", 'wb') as f: pickle.dump(data, f) data["infos"] = val_nusc_infos with open( output_path / f"infos_val_multisweep_{radar_version}_sweeps{max_sweeps}.pkl", 'wb') as f: pickle.dump(data, f)
def __init__(self, config: dict = {}, helper: PredictHelper = None, py_logger=None, tb_logger=None): self.na_config = Configuration({ 'debug': False, 'pred_horizon': 6, 'obs_horizon': 2, 'freq': 2, 'load_dataset': False, 'version': 'v1.0-mini', 'debug': False, 'py_logging_path': None, 'tb_logging_path': None }) self.na_config.update(config) self.name = None self.py_logger = py_logger self.tb_logger = tb_logger self.dataroot = None if 'mini' in self.na_config['version']: self.dataroot = mini_path else: self.dataroot = full_path if self.py_logger is None and self.na_config[ 'py_logging_path'] is not None: print(f"py logging path: {self.na_config['py_logging_path']}") self.py_logger = logger self.py_logger.add(self.na_config['py_logging_path'] + "/log.txt") # self.py_logger = logging.getLogger(self.name) # print(f"py logging path: {self.na_config['py_logging_path']}") # self.py_logger.addHandler(logging.FileHandler(os.path.join(self.na_config['py_logging_path'], 'log.txt'), mode='w')) # if self.py_logger is not None: # self.py_logger.propagate = False if self.tb_logger is None and self.na_config[ 'tb_logging_path'] is not None: self.tb_logger = SummaryWriter( log_dir=os.path.join(self.na_config['tb_logging_path'])) self.helper = helper self.nusc = None if self.helper is not None: self.nusc = self.helper.data else: if self.dataroot is not None and self.na_config[ 'version'] is not None and self.na_config[ 'load_dataset'] and self.helper is None: self.nusc = NuScenes(dataroot=self.dataroot, version=self.na_config['version'], verbose=True) self.helper = PredictHelper(self.nusc) #### Initialize Map #### self.nusc_map_dict = { 'boston-seaport': NuScenesMap(dataroot=self.dataroot, map_name='boston-seaport'), 'singapore-hollandvillage': NuScenesMap(dataroot=self.dataroot, map_name='singapore-hollandvillage'), 'singapore-onenorth': NuScenesMap(dataroot=self.dataroot, map_name='singapore-onenorth'), 'singapore-queenstown': NuScenesMap(dataroot=self.dataroot, map_name='singapore-queenstown'), } #### Initialize CAN #### self.nusc_can = NuScenesCanBus(dataroot=self.dataroot) #### self.all_info = {'config': self.na_config}
class ConvertScene(object): def __init__(self, config={}): self.config = { 'version': 'v1.0-mini', 'dataroot': os.path.join(os.environ['PKG_PATH'],'data') } self.config.update(config) self.nusc = NuScenes(version=self.config['version'], dataroot=self.config['dataroot'], verbose=True) self.nusc_can = NuScenesCanBus(dataroot=self.config['dataroot']) self.utils = Utils(self.nusc, self.nusc_can) def load_experiment_rollout_data(self, experiment_path: str): experiment_result_dict = { 'figures':{ 'bev/bev': {}, 'ctrl/ctrl': {}, 'input_images/current_raster':{}, 'q_transitions/q_transitions':{}, 'agn_node_cams/agn_node_cams':{} }, 'sim_ego_pos': {}, 'sim_ego_quat': {} } #### figures #### for img_folder_name in experiment_result_dict['figures'].keys(): p_list = np.array([str(p) for p in Path(experiment_path+'/'+img_folder_name).glob('*.png')]) idx = np.argsort([int(p.split('/')[-1][:2]) for p in p_list]) p_list = p_list[idx].tolist() for p in p_list: with open(p, 'rb') as png_file: png = png_file.read() idx = int(p.split('/')[-1][:2]) experiment_result_dict['figures'][img_folder_name][idx] = png #### ego #### d = cloudpickle.load(open(os.path.join(experiment_path, f'rollout.pkl'), 'rb')) sim_ego_pos = {} sim_ego_quat = {} for obs in d['obs']: sim_ego_pos[obs['sample_idx']] = obs['sim_ego_pos_gb'].tolist() + [0.01] sim_ego_quat[obs['sample_idx']] = obs['sim_ego_quat_gb'].tolist() experiment_result_dict['sim_ego_pos'] = sim_ego_pos experiment_result_dict['sim_ego_quat'] = sim_ego_quat return experiment_result_dict def load_experiment_rollout_data_pandas(self, experiment_path: str): experiment_result_dict = { 'figures':{ 'feature': {}, 'ioc_param': {}, 'monitor':{} }, 'sim_ego_pos': {}, 'sim_ego_quat': {} } #### figures #### for img_folder_name in experiment_result_dict['figures'].keys(): p_list = np.array([str(p) for p in Path(experiment_path+'/'+img_folder_name).glob('*.png')]) idx = np.argsort([int(p.split('/')[-1][:2]) for p in p_list]) p_list = p_list[idx].tolist() for p in p_list: with open(p, 'rb') as png_file: png = png_file.read() idx = int(p.split('/')[-1][:2]) experiment_result_dict['figures'][img_folder_name][idx] = png #### ego #### d_pd = pd.read_pickle(os.path.join(experiment_path, f'rollout.pkl')) d = d_pd['env_info'] sim_ego_pos = {} sim_ego_quat = {} for obs in d['obs']: sim_ego_pos[obs['sample_idx']] = obs['sim_ego_pos_gb'].tolist() + [0.01] sim_ego_quat[obs['sample_idx']] = obs['sim_ego_quat_gb'].tolist() experiment_result_dict['sim_ego_pos'] = sim_ego_pos experiment_result_dict['sim_ego_quat'] = sim_ego_quat return experiment_result_dict def convert(self, scene_data_p, scene_name=None, model_name=None, up_to_step=None, ckpt_name=None): self.scene_name = scene_name self.model_name = model_name self.ckpt_name = ckpt_name self.scene_data_p = scene_data_p self.experiment_result_dict = self.load_experiment_rollout_data(self.scene_data_p) #self.experiment_result_dict = self.load_experiment_rollout_data_pandas(self.scene_data_p) self.convert_scene(scene_name, self.utils, self.config['dataroot'], up_to_step=up_to_step, experiment_result_dict=self.experiment_result_dict ) def convert_scene(self, scene_name, utils, dataroot, up_to_step, *args, **kwargs): experiment_result_dict = None if 'experiment_result_dict' in kwargs.keys(): experiment_result_dict = kwargs['experiment_result_dict'] scene_i = None for scene_i in self.nusc.scene: if scene_i['name'] == scene_name: scene = scene_i break scene_name = scene['name'] log = self.nusc.get('log', scene['log_token']) location = log['location'] print(f'Loading map "{location}"') self.nusc_map = NuScenesMap(dataroot=dataroot, map_name=location) print(f'Loading bitmap "{self.nusc_map.map_name}"') bitmap = BitMap(self.nusc_map.dataroot, self.nusc_map.map_name, 'basemap') print(f'Loaded {bitmap.image.shape} bitmap') cur_sample = self.nusc.get('sample', scene['first_sample_token']) can_parsers = [ [self.nusc_can.get_messages(scene_name, 'ms_imu'), 0, utils.get_imu_msg], [self.nusc_can.get_messages(scene_name, 'pose'), 0, utils.get_odom_msg], [self.nusc_can.get_messages(scene_name, 'steeranglefeedback'), 0, lambda x: utils.get_basic_can_msg('Steering Angle', x)], [self.nusc_can.get_messages(scene_name, 'vehicle_monitor'), 0, lambda x: utils.get_basic_can_msg('Vehicle Monitor', x)], [self.nusc_can.get_messages(scene_name, 'zoesensors'), 0, lambda x: utils.get_basic_can_msg('Zoe Sensors', x)], [self.nusc_can.get_messages(scene_name, 'zoe_veh_info'), 0, lambda x: utils.get_basic_can_msg('Zoe Vehicle Info', x)], ] bag_name = f"{self.model_name}_{self.ckpt_name}_{self.config['scene_name']}.bag" bag_path = os.path.join(self.scene_data_p, bag_name) print(f'Writing to {bag_path}') bag = rosbag.Bag(bag_path, 'w', compression='lz4') stamp = utils.get_time(self.nusc.get('ego_pose', self.nusc.get('sample_data', cur_sample['data']['LIDAR_TOP'])['ego_pose_token'])) map_msg = utils.get_scene_map(scene, self.nusc_map, bitmap, stamp) centerlines_msg = utils.get_centerline_markers(scene, self.nusc_map, stamp) bag.write('/map', map_msg, stamp) bag.write('/semantic_map', centerlines_msg, stamp) last_map_stamp = stamp #idx = list(experiment_result_dict['sim_ego_pos'].keys())[0] idx = 0 while cur_sample is not None: if up_to_step is not None and idx > up_to_step: break sample_lidar = self.nusc.get('sample_data', cur_sample['data']['LIDAR_TOP']) ego_pose = self.nusc.get('ego_pose', sample_lidar['ego_pose_token']) stamp = utils.get_time(ego_pose) #### write experiment_result to bag #### ego_marker = None if experiment_result_dict is not None: # figures # for k,v in experiment_result_dict['figures'].items(): if idx in v.keys(): png = v[idx] msg = CompressedImage() msg.header.frame_id = k msg.header.stamp = stamp msg.format = 'png' msg.data = png bag.write("/"+k+'/image_rect_compressed', msg, stamp) #print(idx) #print(experiment_result_dict['sim_ego_pos'].keys()) if idx in experiment_result_dict['sim_ego_pos'].keys(): ego_pose = { 'translation': experiment_result_dict['sim_ego_pos'][idx], 'rotation': experiment_result_dict['sim_ego_quat'][idx], 'token': ego_pose['token'], 'timestamp': ego_pose['timestamp'] } if idx == list(experiment_result_dict['sim_ego_pos'].keys())[-1]: next_idx = idx break else: next_idx = idx+1 next_sample = utils.nusc.get('sample',cur_sample['next']) sample_lidar = utils.nusc.get('sample_data', next_sample['data']['LIDAR_TOP']) real_next_ego_pose = utils.nusc.get('ego_pose', sample_lidar['ego_pose_token']) next_ego_pose = { 'translation': experiment_result_dict['sim_ego_pos'][next_idx], 'rotation': experiment_result_dict['sim_ego_quat'][next_idx], 'token': real_next_ego_pose['token'], 'timestamp': real_next_ego_pose['timestamp'] } # publish /tf tf_array = utils.get_tfmessage(cur_sample, current_pose=ego_pose, next_pose=next_ego_pose) bag.write('/tf', tf_array, stamp) else: # publish /tf tf_array = utils.get_tfmessage(cur_sample) bag.write('/tf', tf_array, stamp) else: # publish /tf tf_array = utils.get_tfmessage(cur_sample) bag.write('/tf', tf_array, stamp) # write map topics every two seconds if stamp - rospy.Duration(2.0) >= last_map_stamp: map_msg.header.stamp = stamp for marker in centerlines_msg.markers: marker.header.stamp = stamp bag.write('/map', map_msg, stamp) bag.write('/semantic_map', centerlines_msg, stamp) last_map_stamp = stamp # write CAN messages to /pose, /odom, and /diagnostics can_msg_events = [] for i in range(len(can_parsers)): (can_msgs, index, msg_func) = can_parsers[i] while index < len(can_msgs) and utils.get_utime(can_msgs[index]) < stamp: can_msg_events.append(msg_func(can_msgs[index])) index += 1 can_parsers[i][1] = index can_msg_events.sort(key = lambda x: x[0]) for (msg_stamp, topic, msg) in can_msg_events: bag.write(topic, msg, stamp) # /driveable_area occupancy grid utils.write_occupancy_grid(bag, self.nusc_map, ego_pose, stamp) # iterate sensors for (sensor_id, sample_token) in cur_sample['data'].items(): sample_data = self.nusc.get('sample_data', sample_token) topic = '/' + sensor_id # write the sensor data # if sample_data['sensor_modality'] == 'radar': # msg = utils.get_radar(sample_data, sensor_id) # bag.write(topic, msg, stamp) # elif sample_data['sensor_modality'] == 'lidar': # msg = utils.get_lidar(sample_data, sensor_id) # bag.write(topic, msg, stamp) # if sample_data['sensor_modality'] == 'camera': # msg = utils.get_camera(sample_data, sensor_id, dataroot) # bag.write(topic + '/image_rect_compressed', msg, stamp) # msg = utils.get_camera_info(sample_data, sensor_id) # bag.write(topic + '/camera_info', msg, stamp) # if sample_data['sensor_modality'] == 'camera': # msg = utils.get_lidar_imagemarkers(sample_lidar, sample_data, sensor_id) # bag.write(topic + '/image_markers_lidar', msg, stamp) # utils.write_boxes_imagemarkers(bag, cur_sample['anns'], sample_data, sensor_id, topic, stamp) # publish /markers/annotations marker_array = MarkerArray() ego_marker = Marker() ego_marker.header.frame_id = 'base_link' ego_marker.header.stamp = stamp ego_marker.id = 10000 ego_marker.text = 'ego' ego_marker.type = Marker.CUBE #ego_marker.pose = utils.get_pose(next_ego_pose) ego_marker.frame_locked = True ego_marker.scale.x = 4 ego_marker.scale.y = 1 ego_marker.scale.z = 2 ego_marker.color = utils.make_color((255,61,99), 0.5) marker_array.markers.append(ego_marker) # publish /pose pose_stamped = PoseStamped() pose_stamped.header.frame_id = 'base_link' pose_stamped.header.stamp = stamp pose_stamped.pose.orientation.w = 1 bag.write('/pose', pose_stamped, stamp) # publish /gps coordinates = utils.derive_latlon(location, ego_pose) gps = NavSatFix() gps.header.frame_id = 'base_link' gps.header.stamp = stamp gps.status.status = 1 gps.status.service = 1 gps.latitude = coordinates['latitude'] gps.longitude = coordinates['longitude'] gps.altitude = utils.get_transform(ego_pose).translation.z bag.write('/gps', gps, stamp) for annotation_id in cur_sample['anns']: ann = self.nusc.get('sample_annotation', annotation_id) marker_id = int(ann['instance_token'][:4], 16) c = np.array(self.nusc.explorer.get_color(ann['category_name'])) / 255.0 marker = Marker() marker.header.frame_id = 'map' marker.header.stamp = stamp marker.id = marker_id marker.text = ann['instance_token'][:4] marker.type = Marker.CUBE marker.pose = utils.get_pose(ann) marker.frame_locked = True marker.scale.x = ann['size'][1] marker.scale.y = ann['size'][0] marker.scale.z = ann['size'][2] marker.color = utils.make_color(c, 0.5) marker_array.markers.append(marker) # ego marker # # marker = Marker() # marker.header.frame_id = 'map' # marker.header.stamp = stamp # marker.id = marker_id # marker.text = 'ego' # marker.type = Marker.CUBE # marker.pose = utils.get_pose(ego_pose) # marker.frame_locked = True # marker.scale.x = 4.3 # marker.scale.y = 1.8 # marker.scale.z = 1.6 # c = [0.1,0.1,0.6] # marker.color = utils.make_color(c, 0.9) # marker_array.markers.append(marker) bag.write('/markers/annotations', marker_array, stamp) # collect all sensor frames after this sample but before the next sample non_keyframe_sensor_msgs = [] for (sensor_id, sample_token) in cur_sample['data'].items(): topic = '/' + sensor_id next_sample_token = self.nusc.get('sample_data', sample_token)['next'] while next_sample_token != '': next_sample_data = self.nusc.get('sample_data', next_sample_token) # if next_sample_data['is_key_frame'] or get_time(next_sample_data).to_nsec() > next_stamp.to_nsec(): # break if next_sample_data['is_key_frame']: break # if next_sample_data['sensor_modality'] == 'radar': # msg = utils.get_radar(next_sample_data, sensor_id) # non_keyframe_sensor_msgs.append((msg.header.stamp.to_nsec(), topic, msg)) # elif next_sample_data['sensor_modality'] == 'lidar': # msg = utils.get_lidar(next_sample_data, sensor_id) # non_keyframe_sensor_msgs.append((msg.header.stamp.to_nsec(), topic, msg)) # if next_sample_data['sensor_modality'] == 'camera': # msg = utils.get_camera(next_sample_data, sensor_id, dataroot) # camera_stamp_nsec = msg.header.stamp.to_nsec() # non_keyframe_sensor_msgs.append((camera_stamp_nsec, topic + '/image_rect_compressed', msg)) # msg = utils.get_camera_info(next_sample_data, sensor_id) # non_keyframe_sensor_msgs.append((camera_stamp_nsec, topic + '/camera_info', msg)) # closest_lidar = utils.find_closest_lidar(cur_sample['data']['LIDAR_TOP'], camera_stamp_nsec) # if closest_lidar is not None: # msg = utils.get_lidar_imagemarkers(closest_lidar, next_sample_data, sensor_id) # non_keyframe_sensor_msgs.append((msg.header.stamp.to_nsec(), topic + '/image_markers_lidar', msg)) # else: # msg = utils.get_remove_imagemarkers(sensor_id, 'LIDAR_TOP', msg.header.stamp) # non_keyframe_sensor_msgs.append((msg.header.stamp.to_nsec(), topic + '/image_markers_lidar', msg)) # Delete all image markers on non-keyframe camera images # msg = get_remove_imagemarkers(sensor_id, 'LIDAR_TOP', msg.header.stamp) # non_keyframe_sensor_msgs.append((camera_stamp_nsec, topic + '/image_markers_lidar', msg)) # msg = get_remove_imagemarkers(sensor_id, 'annotations', msg.header.stamp) # non_keyframe_sensor_msgs.append((camera_stamp_nsec, topic + '/image_markers_annotations', msg)) next_sample_token = next_sample_data['next'] # sort and publish the non-keyframe sensor msgs non_keyframe_sensor_msgs.sort(key=lambda x: x[0]) for (_, topic, msg) in non_keyframe_sensor_msgs: bag.write(topic, msg, msg.header.stamp) # move to the next sample cur_sample = self.nusc.get('sample', cur_sample['next']) if cur_sample.get('next') != '' else None idx += 1 bag.close() print(f'Finished writing {bag_name}')
trajectories = np.select(conditions, choices) if debug: trajectories = trajs # return trajectories, t_selections return trajectories if __name__ == "__main__": from nuscenes.nuscenes import NuScenes from nuscenes.can_bus.can_bus_api import NuScenesCanBus import matplotlib.pyplot as plt nusc = NuScenes("v1.0-mini", "/data/nuscenes") nusc_can = NuScenesCanBus(dataroot="/data/nuscenes") for scene in nusc.scene: scene_name = scene["name"] scene_id = int(scene_name[-4:]) if scene_id in nusc_can.can_blacklist: print(f"skipping {scene_name}") continue pose = nusc_can.get_messages(scene_name, "pose") saf = nusc_can.get_messages(scene_name, "steeranglefeedback") vm = nusc_can.get_messages(scene_name, "vehicle_monitor") # NOTE: I tried to verify if the relevant measurements are consistent # across multiple tables that contain redundant information # NOTE: verified pose's velocity matches vehicle monitor's # but the pose table offers at a much higher frequency # NOTE: same that steeranglefeedback's steering angle matches vehicle monitor's