class NuscenesDataset(Dataset):
def __init__(self, data_root, version = 'v1.0-trainval', n: tuple=None):
"""
Args:
root_dir (string): root NuScenes directory
transform (callable, optional): Optional transform to be applied on a sample.
n - tuple with left and right index boundaries, in case we don't want to use all data
"""
self.data_root = data_root
self.nuscenes = NuScenes(version=version, dataroot=data_root, verbose=True)
if n:
assert len(n) == 2
self.samples = self.nuscenes.sample[n[0]:n[1]]
else:
self.samples = self.nuscenes.sample
# point_clouds_features will be of shape (N, M, 5),
# where N - number of samples, M - number of points in according sample's pointcloud
self.point_clouds_features = []
self.point_clouds_labels = []
self.__set_point_clouds()
def __set_point_clouds(self):
for sample in self.samples:
sample_data_token = sample['data']['LIDAR_TOP']
my_sample_lidar_data = self.nuscenes.get('sample_data', sample_data_token)
lidarseg_labels_filename = join(self.data_root,
self.nuscenes.get('lidarseg', sample_data_token)['filename'])
# loading directly from files to perceive the ring_index information
points_raw = np.fromfile(self.data_root + my_sample_lidar_data["filename"], dtype=np.float32).reshape((-1, 5))
point_labels = np.fromfile(lidarseg_labels_filename, dtype=np.uint8)
self.point_clouds_features.append(points_raw)
self.point_clouds_labels.append(point_labels)
self.point_clouds_features = np.array(self.point_clouds_features)
self.point_clouds_labels = np.array(self.point_clouds_labels)
def __len__(self):
return len(self.samples)
def get_front_bb(self, sample: dict):
"""
Function computes and returns
An array of points of a bounding box in sensor coordinates.
Which are in the front 90 degrees
Each point is a (x, y) coordinate, each BB is 4 points
:param sample: nuscenes sample dictionary
:return: np.array of shape (N, 8, 3)
"""
my_sample_lidar_data = self.nuscenes.get('sample_data', sample['data']['LIDAR_TOP'])
sample_annotation = self.nuscenes.get_boxes(my_sample_lidar_data['token'])
ego_record = self.nuscenes.get('ego_pose', my_sample_lidar_data['ego_pose_token'])
cs_record = self.nuscenes.get('calibrated_sensor', my_sample_lidar_data['calibrated_sensor_token'])
# first step: transform from absolute to ego
ego_translation = -np.array(ego_record['translation'])
# second step: transform from ego to sensor
cs_translation = -np.array(cs_record['translation'])
corners = []
for box in sample_annotation:
box.translate(ego_translation)
box.rotate(Quaternion(ego_record['rotation']).inverse)
box.translate(cs_translation)
box.rotate(Quaternion(cs_record['rotation']).inverse)
# at this point bounding boxes are in sensor coordinate system
# now we want to exclude such BB that do not have their center
# lying in the front 90 degrees
if box.center[1] <= 0:
continue
box.center_azimuth = np.degrees(np.arctan(box.center[0] / box.center[1]))
# Transform front azimuth to be in range from 0 to 180
box.center_azimuth = 90 - box.center_azimuth
if not (45 < box.center_azimuth < 135):
continue
corners.append(box.bottom_corners())
if len(corners) > 0:
return np.transpose(np.array(corners), (0, 2, 1))
else:
return np.zeros((1, 4, 3))
def points_in_box(self, coordinates, bounding_box_corners):
"""
bounding_box_corners: bbc of a single bb
return a mask of whether points that are in the box
"""
coords_x = coordinates[0]
coords_y = coordinates[1]
min_bb_x = bounding_box_corners[:, 0].min()
max_bb_x = bounding_box_corners[:, 0].max()
min_bb_y = bounding_box_corners[:, 1].min()
max_bb_y = bounding_box_corners[:, 1].max()
c1 = min_bb_x <= coords_x # left_top/left_bottom.x <= coordinate.x
c2 = max_bb_x >= coords_x # right_bottom/right_top.x >= coordinate.x
c3 = min_bb_y <= coords_y # left/right_bottom.y <= coordinate.y
c4 = max_bb_y >= coords_y # right_top/left_top.y >= coordinate.y
c = np.logical_and(np.logical_and(c1, c2),
np.logical_and(c3, c4))
return c
def get_bb_targets(self, idx, bounding_box_corners):
coordinates = self.point_clouds_features[idx][:2]
for bb_c in bounding_box_corners:
if self.points_in_box(coordinates, bb_c).any():
return np.array(bb_c[:, :2])
else:
return np.zeros_like(bb_c[:, :2])
def __getitem__(self, idx, compute_boxes=True):
"""
compute_box will be set to False in child classes, so target boxes are computed after rotations
"""
sample_idx = self.samples[idx]
pcl_features = self.point_clouds_features[idx]
pcl_labels = self.point_clouds_labels[idx]
if compute_boxes:
front_bbs = self.get_front_bb(sample_idx)
target_bounding_boxes = self.get_bb_targets(idx, front_bbs)
return pcl_features, pcl_labels, target_bounding_boxes
else:
return pcl_features, pcl_labels, []
class CustomNuscDatasetMMS(Dataset):
NumPointFeatures = 4
def __init__(
self,
root_path=f'/media/starlet/LdTho/data/sets/nuscenes/v1.0-{VERSION}',
info_path=None,
class_names=["traffic_cone"],
prep_func=None,
num_point_features=None):
self.NumPointFeatures = 4
self.class_names = class_names
self.nusc = NuScenes(dataroot=root_path, version=f'v1.0-{VERSION}')
self._prep_func = prep_func
self.filtered_sample_tokens = []
for sample in self.nusc.sample:
sample_token = sample['token']
sample_lidar_token = sample['data']['LIDAR_TOP']
boxes = self.nusc.get_boxes(sample_lidar_token)
box_names = [
NameMapping[b.name] for b in boxes
if b.name in NameMapping.keys()
]
for box in boxes:
if box.name not in NameMapping.keys():
continue
# if NameMapping[box.name] in self.class_names:
if (NameMapping[box.name] in [
"traffic_cone"
]) & (box_names.count('traffic_cone') > MIN_CONES_PER_SAMPLE):
self.filtered_sample_tokens.append(sample_token)
break
self.filtered_sample_tokens = self.filtered_sample_tokens[:round(
len(self.filtered_sample_tokens) * TRAINVAL_SPLIT_PERCENTAGE)]
self.split = np.arange(len(self.filtered_sample_tokens))
def __len__(self):
return self.split.shape[0]
def __getitem__(self, index):
input_dict = self.get_sensor_data(index)
example = self._prep_func(input_dict=input_dict)
example["metadata"] = input_dict["metadata"]
if "anchors_mask" in example:
example["anchors_mask"] = example["anchors_mask"].astype(np.uint8)
return example
def get_sensor_data(self, query, token=None):
res = {
'lidar': {
'type': 'lidar',
'points': None,
},
'metadata': {
'token': self.filtered_sample_tokens[query]
}
}
if token:
query = self.filtered_sample_tokens.index(token)
points = self.getPoints(query)
boxes_dict = self.getBoxes(query)
res['lidar']['points'] = points
gt_boxes = []
gt_names = []
for box in boxes_dict:
xyz = box.center
wlh = box.wlh
theta = quaternion_yaw(box.orientation)
gt_boxes.append([
xyz[0], xyz[1], xyz[2], wlh[0], wlh[1], wlh[2],
-theta - np.pi / 2
])
gt_names.append(box.name)
gt_boxes = np.concatenate(gt_boxes).reshape(-1, 7)
gt_names = np.array(gt_names)
res['lidar']['annotations'] = {
'boxes': gt_boxes,
'names': gt_names,
}
return res
###
def getPoints(self, index):
sample = self.nusc.get('sample', self.filtered_sample_tokens[index])
sample_lidar_token = sample['data']['LIDAR_TOP']
lidar_data = self.nusc.get('sample_data', sample_lidar_token)
ego_pose = self.nusc.get('ego_pose', lidar_data['ego_pose_token'])
calibrated_sensor = self.nusc.get(
'calibrated_sensor', lidar_data['calibrated_sensor_token'])
global_from_car = transform_matrix(ego_pose['translation'],
Quaternion(ego_pose['rotation']),
inverse=False)
car_from_sensor = transform_matrix(calibrated_sensor['translation'],
Quaternion(
calibrated_sensor['rotation']),
inverse=False)
try:
lidar_pointcloud, times = LidarPointCloud.from_file_multisweep(
self.nusc, sample, 'LIDAR_TOP', 'LIDAR_TOP')
lidar_pointcloud.transform(car_from_sensor)
except Exception as e:
print(f"Failed to load Lidar Pointcloud for {sample}:{e}")
points = lidar_pointcloud.points
points[3, :] /= 255
points[3, :] -= 0.5
# points_cat = np.concatenate([points, times], axis=0).transpose()
points_cat = points.transpose()
points_cat = points_cat[~np.isnan(points_cat).any(axis=1)]
return points_cat
def getBoxes(self, index):
sample = self.nusc.get('sample', self.filtered_sample_tokens[index])
sample_lidar_token = sample['data']['LIDAR_TOP']
lidar_data = self.nusc.get('sample_data', sample_lidar_token)
ego_pose = self.nusc.get('ego_pose', lidar_data['ego_pose_token'])
boxes_dict = self.nusc.get_boxes(sample_lidar_token)
keep_box_idx = []
for i, box in enumerate(boxes_dict):
if box.name not in NameMapping.keys():
continue
if NameMapping[box.name] in self.class_names:
box.name = NameMapping[box.name]
keep_box_idx.append(i)
boxes_dict = [
box for i, box in enumerate(boxes_dict) if i in keep_box_idx
]
self.move_boxes_to_car_space(boxes_dict, ego_pose)
# print(boxes_dict)
return boxes_dict
def move_boxes_to_car_space(self, boxes, ego_pose):
"""
Move boxes from world space to car space.
Note: mutates input boxes.
"""
translation = -np.array(ego_pose['translation'])
rotation = Quaternion(ego_pose['rotation']).inverse
for box in boxes:
# Bring box to car space
box.translate(translation)
box.rotate(rotation)
def parse_sequence(
data: NuScenes, add_nonkey_frames: bool,
scene_info: Tuple[str, str]) -> Tuple[List[Frame], List[FrameGroup]]:
"""Parse a full NuScenes sequence and convert it into scalabel frames."""
sample_token, scene_name = scene_info
frames, groups = [], []
frame_index = 0
while sample_token:
sample = data.get("sample", sample_token)
lidar_token = sample["data"]["LIDAR_TOP"]
lidar_data = data.get("sample_data", lidar_token)
calibration_lidar = data.get("calibrated_sensor",
lidar_data["calibrated_sensor_token"])
timestamp = lidar_data["timestamp"]
lidar_filepath = lidar_data["filename"]
ego_pose = data.get("ego_pose", lidar_data["ego_pose_token"])
frame_names = []
next_nonkey_frames = []
for cam in cams:
cam_token = sample["data"][cam]
boxes = data.get_boxes(lidar_token)
frame, next_token = parse_frame(data, scene_name, frame_index,
cam_token, boxes)
frame_names.append(frame.name)
frames.append(frame)
if add_nonkey_frames and next_token is not None:
next_nonkey_frames.append(next_token)
group = FrameGroup(
name=sample_token,
videoName=scene_name,
frameIndex=frame_index,
url=lidar_filepath,
extrinsics=get_extrinsics(ego_pose, calibration_lidar),
timestamp=timestamp,
frames=frame_names,
labels=parse_labels(data, data.get_boxes(lidar_token), ego_pose,
calibration_lidar),
)
groups.append(group)
frame_index += 1
nonkey_count = 0
while len(next_nonkey_frames) > 0:
new_next_nonkey_frames = []
frame_names = []
nonkey_count += 1
for cam_token in next_nonkey_frames:
assert cam_token is not None, "camera for non-key missing!"
frame, next_token = parse_frame(data, scene_name, frame_index,
cam_token)
frame_names.append(frame.name)
frames.append(frame)
if add_nonkey_frames and next_token is not None:
new_next_nonkey_frames.append(next_token)
group = FrameGroup(
name=sample_token + f"_{nonkey_count}",
videoName=scene_name,
frameIndex=frame_index,
frames=frame_names,
)
groups.append(group)
next_nonkey_frames = new_next_nonkey_frames
frame_index += 1
sample_token = sample["next"]
return frames, groups
class CustomNuscEvalDataset(Dataset):
NumPointFeatures = 5
def __init__(
self,
root_path=f'/media/starlet/LdTho/data/sets/nuscenes/v1.0-{EVAL_VERSION}',
info_path=None,
class_names=["traffic_cone"],
prep_func=None,
num_point_features=None):
self.NumPointFeatures = 5
self.class_names = class_names
self.nusc = NuScenes(dataroot=root_path,
version=f'v1.0-{EVAL_VERSION}')
self._prep_func = prep_func
self.root_path = root_path
self.eval_version = "detection_cvpr_2019"
# self.filtered_sample_tokens = [s['token'] for s in self.nusc.sample]
self.filtered_sample_tokens = []
for sample in self.nusc.sample:
sample_token = sample['token']
sample_lidar_token = sample['data']['LIDAR_TOP']
boxes = self.nusc.get_boxes(sample_lidar_token)
box_names = [
NameMapping[b.name] for b in boxes
if b.name in NameMapping.keys()
]
for box in boxes:
if box.name not in NameMapping.keys():
continue
# if NameMapping[box.name] in self.class_names:
if (NameMapping[box.name] in [
"traffic_cone"
]) & (box_names.count('traffic_cone') > MIN_CONES_PER_SAMPLE):
self.filtered_sample_tokens.append(sample_token)
break
if EVAL_VERSION == "trainval":
self.filtered_sample_tokens = self.filtered_sample_tokens[round(
len(self.filtered_sample_tokens) * TRAINVAL_SPLIT_PERCENTAGE):]
self.split = np.arange(len(self.filtered_sample_tokens))
def __len__(self):
return self.split.shape[0]
def __getitem__(self, index):
input_dict = self.get_sensor_data(index)
example = self._prep_func(input_dict=input_dict)
example["metadata"] = input_dict["metadata"]
if "anchors_mask" in example:
example["anchors_mask"] = example["anchors_mask"].astype(np.uint8)
return example
def get_sensor_data(self, query, token=None):
res = {
'lidar': {
'type': 'lidar',
'points': None,
},
'metadata': {
'token': self.filtered_sample_tokens[query]
}
}
if token:
query = self.filtered_sample_tokens.index(token)
points = self.getPoints(query)
boxes_dict = self.getBoxes(query)
res['lidar']['points'] = points
gt_boxes = []
gt_names = []
for box in boxes_dict:
xyz = box.center
wlh = box.wlh
theta = quaternion_yaw(box.orientation)
gt_boxes.append([
xyz[0], xyz[1], xyz[2], wlh[0], wlh[1], wlh[2],
-theta - np.pi / 2
])
gt_names.append(box.name)
gt_boxes = np.concatenate(gt_boxes).reshape(-1, 7)
gt_names = np.array(gt_names)
res['lidar']['annotations'] = {
'boxes': gt_boxes,
'names': gt_names,
}
return res
###
def getPoints(self, index):
sample = self.nusc.get('sample', self.filtered_sample_tokens[index])
sample_lidar_token = sample['data']['LIDAR_TOP']
lidar_data = self.nusc.get('sample_data', sample_lidar_token)
ego_pose = self.nusc.get('ego_pose', lidar_data['ego_pose_token'])
calibrated_sensor = self.nusc.get(
'calibrated_sensor', lidar_data['calibrated_sensor_token'])
global_from_car = transform_matrix(ego_pose['translation'],
Quaternion(ego_pose['rotation']),
inverse=False)
car_from_sensor = transform_matrix(calibrated_sensor['translation'],
Quaternion(
calibrated_sensor['rotation']),
inverse=False)
try:
lidar_pointcloud, times = LidarPointCloud.from_file_multisweep(
self.nusc, sample, 'LIDAR_TOP', 'LIDAR_TOP')
lidar_pointcloud.transform(car_from_sensor)
except Exception as e:
print(f"Failed to load Lidar Pointcloud for {sample}:{e}")
points = lidar_pointcloud.points
points[3, :] /= 255
points[3, :] -= 0.5
points_cat = np.concatenate([points, times], axis=0).transpose()
points_cat = points_cat[~np.isnan(points_cat).any(axis=1)]
return points_cat
def getBoxes(self, index):
sample = self.nusc.get('sample', self.filtered_sample_tokens[index])
sample_lidar_token = sample['data']['LIDAR_TOP']
lidar_data = self.nusc.get('sample_data', sample_lidar_token)
ego_pose = self.nusc.get('ego_pose', lidar_data['ego_pose_token'])
boxes_dict = self.nusc.get_boxes(sample_lidar_token)
keep_box_idx = []
for i, box in enumerate(boxes_dict):
if box.name not in NameMapping.keys():
continue
if NameMapping[box.name] in self.class_names:
box.name = NameMapping[box.name]
keep_box_idx.append(i)
boxes_dict = [
box for i, box in enumerate(boxes_dict) if i in keep_box_idx
]
self.move_boxes_to_car_space(boxes_dict, ego_pose)
# print(boxes_dict)
return boxes_dict
def move_boxes_to_car_space(self, boxes, ego_pose, eval=False):
"""
Move boxes from world space to car space.
Note: mutates input boxes.
"""
translation = -np.array(ego_pose['translation'])
rotation = Quaternion(ego_pose['rotation']).inverse
box_list = []
for box in boxes:
# Bring box to car space
box.translate(translation)
box.rotate(rotation)
box_list.append(box)
return box_list
def evaluation(self, detections, output_dir):
"""kitti evaluation is very slow, remove it.
"""
# res_kitti = self.evaluation_kitti(detections, output_dir)
res_nusc = self.evaluation_nusc(detections, output_dir)
res = {
"results": {
"nusc": res_nusc["results"]["nusc"],
# "kitti.official": res_kitti["results"]["official"],
# "kitti.coco": res_kitti["results"]["coco"],
},
"detail": {
"eval.nusc": res_nusc["detail"]["nusc"],
# "eval.kitti": {
# "official": res_kitti["detail"]["official"],
# "coco": res_kitti["detail"]["coco"],
# },
},
}
return res
def evaluation_nusc(self, detections, output_dir):
mapped_class_names = self.class_names
nusc_annos = {}
token2info = {}
for det in detections:
annos = []
boxes = _second_det_to_nusc_box(det)
boxes = self.transform_box_back_to_global(boxes,
det["metadata"]["token"])
# boxes = self.move_boxes_to_car_space(boxes, ego_pose)
for i, box in enumerate(boxes):
name = mapped_class_names[box.label]
velocity = [np.nan, np.nan]
nusc_anno = {
"sample_token": det["metadata"]["token"],
"translation": box.center.tolist(),
"size": box.wlh.tolist(),
"rotation": box.orientation.elements.tolist(),
"velocity": velocity,
"detection_name": name,
"detection_score": box.score,
"attribute_name": DefaultAttribute[name]
}
annos.append(nusc_anno)
nusc_annos[det["metadata"]["token"]] = annos
nusc_submissions = {
"meta": {
"use_camera": False,
"use_lidar": False,
"use_radar": False,
"use_map": False,
"use_external": False
},
"results": nusc_annos
}
res_path = Path(output_dir) / "result_nusc.json"
with open(res_path, "w") as f:
json.dump(nusc_submissions, f)
eval_main_file = Path(__file__).resolve().parent / "nusc_eval.py"
cmd = f"python {str(eval_main_file)} --root_path=\"{self.root_path}\""
cmd += f" --version=\"v1.0-mini\" --eval_version={self.eval_version}"
cmd += f" --res_path=\"{str(res_path)}\" --eval_set=mini_train"
cmd += f" --output_dir=\"{output_dir}\""
print(cmd)
subprocess.check_output(cmd, shell=True)
with open(Path(output_dir) / "metrics_summary.json", "r") as f:
metrics = json.load(f)
detail = {}
res_path.unlink()
result = f"Nusc {VERSION} evaluation"
for name in mapped_class_names:
detail[name] = {}
for k, v in metrics["label_aps"][name].items():
detail[name][f"dist@{k}"] = v
tp_errs = []
tp_names = []
for k, v in metrics["label_tp_errors"][name].items():
detail[name][k] = v
tp_errs.append(f"{v:.4f}")
tp_names.append(k)
threshs = ', '.join(list(metrics["label_aps"][name].keys()))
scores = list(metrics["label_aps"][name].values())
scores = ', '.join([f"{s * 100:.2f}" for s in scores])
result += f"{name} Nusc dist AP@{threshs} and TP errors\n"
result += scores
result += ', '.join(tp_names) + ": " + ", ".join(tp_errs)
result += "\n"
return {
"results": {
"nusc": result,
},
"detail": {
"nusc": detail
}
}
def transform_box_back_to_global(self, boxes, token):
sample = self.nusc.get('sample', token)
sample_lidar_token = sample['data']['LIDAR_TOP']
lidar_data = self.nusc.get('sample_data', sample_lidar_token)
ego_pose = self.nusc.get('ego_pose', lidar_data['ego_pose_token'])
translation = np.array(ego_pose['translation'])
rotation = Quaternion(ego_pose['rotation'])
box_list = []
for box in boxes:
box.rotate(rotation)
box.translate(translation)
box_list.append(box)
return box_list