def make_anchor_boxes():
"""
constructs the anchor boxes for the model, returns a list of
Box objects, an array of xy coordinates of the corners of shape
[N,4,2], an array of the centers of shape [N,3] and an array of
the boxes in (x1,y1,x2,y2) format where (x1,y1) are the top left
corner and (x2,y2) the bottom right
"""
fm_height = cfg.DATA.FM_HEIGHT
fm_width = cfg.DATA.FM_WIDTH
fm_scale = cfg.DATA.FM_SCALE
anchor_dims = cfg.DATA.ANCHOR_DIMS
anchor_yaws = cfg.DATA.ANCHOR_YAWS
anchor_zs = cfg.DATA.ANCHOR_ZS
corners_list = []
boxes_list = []
centers_list = []
xy_list = []
# loop through the height,width and dimensions of the boxes
for y in tqdm(range(0, fm_height)):
for x in range(0, fm_width):
for d in range(0, len(anchor_dims)):
x_center = (x + 0.5) / fm_scale
y_center = (y + 0.5) / fm_scale
# the wlh, z coord and yaw of the anchors are set in config
z_center = anchor_zs[d]
width = anchor_dims[d][0]
length = anchor_dims[d][1]
height = anchor_dims[d][2]
yaw = anchor_yaws[d]
quat = Quaternion(axis=[0, 0, 1], degrees=yaw)
# create new box object
box = Box(center=[x_center, y_center, z_center],
size=[width, length, height],
orientation=quat)
boxes_list.append(box)
# take only the xy coords of the corners
bc = box.bottom_corners().transpose([1, 0])
corners_list.append(bc[:, :2])
centers_list.append([x_center, y_center, z_center])
if yaw > 0:
xy_list.append(np.concatenate((bc[1, :2], bc[3, :2])))
else:
xy_list.append(np.concatenate((bc[2, :2], bc[0, :2])))
return boxes_list, np.array(corners_list), np.array(
centers_list), np.array(xy_list)
def test_rotate_around_box_center(size, center, translate, angle1, angle2):
axis = [0, 0, 1]
q1 = Quaternion(axis=axis, angle=angle1)
q2 = Quaternion(axis=axis, angle=angle2)
minus_q2 = Quaternion(axis=axis, angle=-q2.angle)
original = Box(center=center, size=size, orientation=q1)
assert original == (original.copy().rotate_around_box_center(
q2).rotate_around_box_center(minus_q2))
assert original == (original.copy().rotate_around_box_center(q2).translate(
translate).rotate_around_box_center(minus_q2).translate(-translate))
def transform_box_from_world_to_flat_sensor_coordinates(first_train_sample_box: Box, sample_data_token: str,
lyftd: LyftDataset):
sample_box = first_train_sample_box.copy()
sd_record = lyftd.get("sample_data", sample_data_token)
cs_record = lyftd.get("calibrated_sensor", sd_record["calibrated_sensor_token"])
sensor_record = lyftd.get("sensor", cs_record["sensor_token"])
pose_record = lyftd.get("ego_pose", sd_record["ego_pose_token"])
# Move box to ego vehicle coord system parallel to world z plane
ypr = Quaternion(pose_record["rotation"]).yaw_pitch_roll
yaw = ypr[0]
sample_box.translate(-np.array(pose_record["translation"]))
sample_box.rotate(Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)]).inverse)
# Move box to sensor vehicle coord system parallel to world z plane
# We need to steps, because camera coordinate is x(right), z(front), y(down)
inv_ypr = Quaternion(cs_record["rotation"]).inverse.yaw_pitch_roll
angz = inv_ypr[0]
angx = inv_ypr[2]
sample_box.translate(-np.array(cs_record['translation']))
# rotate around z-axis
sample_box.rotate(Quaternion(scalar=np.cos(angz / 2), vector=[0, 0, np.sin(angz / 2)]))
# rotate around x-axis (by 90 degrees)
angx = 90
sample_box.rotate(Quaternion(scalar=np.cos(angx / 2), vector=[np.sin(angx / 2), 0, 0]))
return sample_box
def get_pred_str(pred, sample_token):
boxes_lidar = pred["box3d_lidar"]
boxes_class = pred["label_preds"]
scores = pred['scores']
preds_classes = [classes[x] for x in boxes_class]
box_centers = boxes_lidar[:, :3]
box_yaws = boxes_lidar[:, -1]
box_wlh = boxes_lidar[:, 3:6]
info = token2info[sample_token] # a `sample` token
boxes = []
pred_str = ''
for idx in range(len(boxes_lidar)):
translation = box_centers[idx]
yaw = -box_yaws[idx] - pi / 2
size = box_wlh[idx]
name = preds_classes[idx]
detection_score = scores[idx]
quat = Quaternion(scalar=np.cos(yaw / 2),
vector=[0, 0, np.sin(yaw / 2)])
box = Box(center=box_centers[idx],
size=size,
orientation=quat,
score=detection_score,
name=name,
token=sample_token)
box = to_glb(box, info)
pred = str(box.score) + ' ' + str(box.center[0]) + ' ' \
+ str(box.center[1]) + ' ' + str(box.center[2]) + ' ' \
+ str(box.wlh[0]) + ' ' + str(box.wlh[1]) + ' ' + \
str(box.wlh[2]) + ' ' + str(box.orientation.yaw_pitch_roll[0]) \
+ ' ' + str(name) + ' '
pred_str += pred
return pred_str.strip()
def get_box_from_inference(lyftd: LyftDataset, heading_cls, heading_res,
rot_angle, size_cls, size_res, center_coord,
sample_data_token, score, type_name: str) -> Box:
heading_angle = get_heading_angle(heading_cls, heading_res, rot_angle)
size = get_size(size_cls, size_res)
rot_angle += np.pi / 2
center_sensor_coord = get_center_in_sensor_coord(center_coord=center_coord,
rot_angle=rot_angle)
# Make Box
# The rationale of doing this: to conform the convention of Box class, the car is originally heading to +x axis,
# with y(left) and z(top). To make the car heading to the angle it should be in the camera coordinate,
# we have to rotate it by 90 degree around x axis and [theta] degree around y axis, where [theta] is the heading angle
l, w, h = size
first_rot = Quaternion(axis=[1, 0, 0], angle=np.pi / 2)
second_rot = Quaternion(axis=[0, -1, 0], angle=-heading_angle)
box_in_sensor_coord = Box(center=center_sensor_coord,
size=[w, l, h],
orientation=second_rot * first_rot,
score=score,
name=type_name)
box_in_world_coord = convert_box_to_world_coord_with_sample_data_token(
box_in_sensor_coord, sample_data_token, lyftd)
# assert np.abs(box_in_world_coord.orientation.axis[0]) <= 0.02
# assert np.abs(box_in_world_coord.orientation.axis[1]) <= 0.02
return box_in_world_coord
def _second_det_to_nusc_box(detection):
from lyft_dataset_sdk.utils.data_classes import Box
box3d = detection["box3d_lidar"].cpu().numpy()
scores = detection["scores"].cpu().numpy()
labels = detection["label_preds"].cpu().numpy()
box3d[:, 6] = -box3d[:, 6] - np.pi / 2
box_list = []
for i in range(box3d.shape[0]):
quat = Quaternion(axis=[0, 0, 1], radians=box3d[i, 6])
velocity = (np.nan, np.nan, np.nan)
if box3d.shape[1] == 9:
velocity = (*box3d[i, 7:9], 0.0)
# velo_val = np.linalg.norm(box3d[i, 7:9])
# velo_ori = box3d[i, 6]
# velocity = (velo_val * np.cos(velo_ori), velo_val * np.sin(velo_ori), 0.0)
box = Box(box3d[i, :3],
box3d[i, 3:6],
quat,
label=labels[i],
score=scores[i],
velocity=velocity)
box_list.append(box)
return box_list
def move_box_to_car_space(box, image=True):
"""
takes a box in image space and moves it to the space
of the ego vehicle
"""
#pdb.set_trace()
x, y, z = box.center
w, l, h = box.wlh
canv_h = cfg.DATA.CANVAS_HEIGHT
y_step = cfg.DATA.Y_STEP
x_step = cfg.DATA.X_STEP
y_min = cfg.DATA.Y_MIN
x_min = cfg.DATA.X_MIN
if image:
y = (canv_h - 1) - y
car_y = y * y_step + y_min
car_x = x * x_step + x_min
car_w = w * y_step
car_l = l * x_step
box_wlh = np.array([car_w, car_l, h])
box_center = np.array([car_x, car_y, z])
car_box = Box(center=box_center,
size=box_wlh,
orientation=Quaternion(list(box.orientation)),
name=box.name,
token=box.token,
score=box.score)
return car_box
def move_boxes_to_canvas_space(boxes, ego_pose, lidar_points):
"""
takes a list of ground truth boxes in global space
and moves them to canvas space. We first move the
boxes to the ego car coordinate system, then we
move the boxes to the voxelized canvas space
"""
box_list = []
x_min = cfg.DATA.X_MIN
x_max = cfg.DATA.X_MAX
y_min = cfg.DATA.Y_MIN
y_max = cfg.DATA.Y_MAX
z_min = cfg.DATA.Z_MIN
z_max = cfg.DATA.Z_MAX
x_step = cfg.DATA.X_STEP
y_step = cfg.DATA.Y_STEP
box_translation = -np.array(ego_pose['translation'])
box_rotation = Quaternion(ego_pose['rotation']).inverse
for box in boxes:
#transform to car space
box.translate(box_translation)
box.rotate(box_rotation)
# ignore boxes that are outside of the bounds
# of the lidar point cloud
box_x, box_y, box_z = box.center
if (box_x < x_min) or (box_x > x_max) or \
(box_y < y_min) or (box_y > y_max) or \
(box_z < z_min) or (box_z > z_max):
continue
pts_in_box = np.sum(points_in_box(box, lidar_points))
if pts_in_box < 10: continue
# compute new xy coordinates in canvas space
canv_x = (box_x - cfg.DATA.X_MIN) / x_step
canv_y = (box_y - cfg.DATA.Y_MIN) / y_step
canv_z = box_z
# adjust wlh
box_w, box_l, box_h = box.wlh
canv_w = box_w / y_step
canv_l = box_l / x_step
canv_h = box_h
box_wlh = np.array([canv_w, canv_l, canv_h])
box_center = np.array([canv_x, canv_y, canv_z])
box_list.append(
Box(box_center,
box_wlh,
label=box.label,
orientation=Quaternion(list(box.orientation)),
name=box.name,
token=box.token))
return box_list
def convert_box_to_world_coord(box: Box, sample_token, sensor_type, lyftd: LyftDataset):
sample_box = box.copy()
sample_record = lyftd.get('sample', sample_token)
sample_data_token = sample_record['data'][sensor_type]
converted_sample_box = convert_box_to_world_coord_with_sample_data_token(sample_box, sample_data_token)
return converted_sample_box
def project_kitti_box_to_image(
box: Box, p_left: np.ndarray,
imsize: Tuple[int, int]) -> Union[None, Tuple[int, int, int, int]]:
"""Projects 3D box into KITTI image FOV.
Args:
box: 3D box in KITTI reference frame.
p_left: <np.float: 3, 4>. Projection matrix.
imsize: (width, height). Image size.
Returns: (xmin, ymin, xmax, ymax). Bounding box in image plane or None if box is not in the image.
"""
# Create a new box.
box = box.copy()
# KITTI defines the box center as the bottom center of the object.
# We use the true center, so we need to adjust half height in negative y direction.
box.translate(np.array([0, -box.wlh[2] / 2, 0]))
# Check that some corners are inside the image.
corners = np.array(
[corner for corner in box.corners().T if corner[2] > 0]).T
if len(corners) == 0:
return None
# Project corners that are in front of the camera to 2d to get bbox in pixel coords.
imcorners = view_points(corners, p_left, normalize=True)[:2]
bbox = (np.min(imcorners[0]), np.min(imcorners[1]),
np.max(imcorners[0]), np.max(imcorners[1]))
# Crop bbox to prevent it extending outside image.
bbox_crop = tuple(max(0, b) for b in bbox)
bbox_crop = (
min(imsize[0], bbox_crop[0]),
min(imsize[0], bbox_crop[1]),
min(imsize[0], bbox_crop[2]),
min(imsize[1], bbox_crop[3]),
)
# Detect if a cropped box is empty.
if bbox_crop[0] >= bbox_crop[2] or bbox_crop[1] >= bbox_crop[3]:
return None
return bbox_crop
def get_train_data_sample_token_and_box(idx, train_objects):
first_train_id = train_objects.iloc[idx, 0]
# Make box
orient_q = Quaternion(axis=[0, 0, 1], angle=float(train_objects.loc[idx, 'yaw']))
center_pos = train_objects.iloc[idx, 2:5].values
wlh = train_objects.iloc[idx, 5:8].values
obj_name = train_objects.iloc[idx, -1]
first_train_sample_box = Box(center=list(center_pos), size=list(wlh),
orientation=orient_q, name=obj_name)
return first_train_id, first_train_sample_box
def make_pred_boxes(inds, anchor_box_list, reg, classes, scores, token):
"""
takes a list of positively labeled anchor boxes and adjusts
them using the network adjustment output.
`inds` contains the box indices of the positive anchor boxes
`reg` contains the model outputs for the anchor box adjustments
`classes` contains the classes of the anchors
`scores` contains the model confidences
"""
out_box_list = []
reg = reg.cpu().numpy()
classes = classes.cpu().numpy()
scores = scores.cpu().numpy()
# loop through the positive anchors
for i in inds:
# get the Box object and the box adjustments for box i
a_box = anchor_box_list[i]
offsets = reg[i, :]
a_box_diag = np.sqrt(a_box.wlh[0]**2 + a_box.wlh[1]**2)
# compute new box xyz and wlh
box_x = a_box.center[0] + offsets[0] * a_box_diag
box_y = a_box.center[1] + offsets[1] * a_box_diag
box_z = a_box.center[2] + offsets[2] * a_box.wlh[2]
box_w = np.exp(offsets[3]) * a_box.wlh[0]
box_l = np.exp(offsets[4]) * a_box.wlh[1]
box_h = np.exp(offsets[5]) * a_box.wlh[2]
# get class & yaw of box
box_name = cfg.DATA.IND_TO_NAME[str(int(classes[i]))]
box_yaw = np.arcsin(offsets[6]) + a_box.orientation.yaw_pitch_roll[0]
"""
if sigmoid(offsets[7]) < .5:
box_ort = 1
else:
box_ort = -1
box_yaw *= box_ort
"""
# make new Box object
quat = Quaternion(axis=[0, 0, 1], radians=box_yaw)
box = Box(center=[box_x, box_y, box_z],
size=[box_w, box_l, box_h],
orientation=quat,
name=box_name,
score=scores[i],
token=token)
out_box_list.append(box)
return out_box_list
def transform_box_from_world_to_ego_coordinates(first_train_sample_box: Box, sample_data_token: str,
lyftd: LyftDataset):
sample_box = first_train_sample_box.copy()
sd_record = lyftd.get("sample_data", sample_data_token)
cs_record = lyftd.get("calibrated_sensor", sd_record["calibrated_sensor_token"])
sensor_record = lyftd.get("sensor", cs_record["sensor_token"])
pose_record = lyftd.get("ego_pose", sd_record["ego_pose_token"])
# Move box to ego vehicle coord system
sample_box.translate(-np.array(pose_record["translation"]))
sample_box.rotate(Quaternion(pose_record["rotation"]).inverse)
return sample_box
def get_box_yaw_angle_in_camera_coords(box: Box):
"""
Calculate the heading angle, using the convention in KITTI labels.
:param box: bouding box
:return:
"""
box_corners = box.corners()
v = box_corners[:, 0] - box_corners[:, 4]
heading_angle = np.arctan2(-v[2], v[0])
return heading_angle
def get_box_yaw_angle_in_world_coords(box: Box):
"""
Calculate the heading angle, using world coordinates.
:param box: bouding box
:return:
"""
box_corners = box.corners()
v = box_corners[:, 0] - box_corners[:, 4]
heading_angle = np.arctan2(v[1], v[0])
return heading_angle
def boxes_lidar_to_lyft(boxes3d, scores=None, labels=None):
box_list = []
for k in range(boxes3d.shape[0]):
quat = Quaternion(axis=[0, 0, 1], radians=boxes3d[k, 6])
box = Box(
boxes3d[k, :3],
boxes3d[k, [4, 3, 5]], # wlh
quat,
label=labels[k] if labels is not None else np.nan,
score=scores[k] if scores is not None else np.nan,
)
box_list.append(box)
return box_list
def get_box(sample_annotation_token, lyftdata):
"""Instantiates a Box class from a sample annotation record.
Args:
sample_annotation_token: Unique sample_annotation identifier.
Returns:
"""
record = lyftdata.get("sample_annotation", sample_annotation_token)
return Box(
record["translation"],
record["size"],
Quaternion(record["rotation"]),
name=record["category_name"],
token=record["token"],
)
def test_rotate_around_origin_xy(size, angle1, angle2, center):
x, y, z = center
axis = [0, 0, 1]
q1 = Quaternion(axis=axis, angle=angle1)
q2 = Quaternion(axis=axis, angle=angle2)
minus_q2 = Quaternion(axis=axis, angle=-q2.angle)
original = Box(center=(x, y, z), size=size, orientation=q1)
assert original == (original.copy().rotate_around_box_center(
q2).rotate_around_box_center(minus_q2))
cos_angle2 = q2.rotation_matrix[0, 0]
sin_angle2 = q2.rotation_matrix[1, 0]
new_center = x * cos_angle2 - y * sin_angle2, x * sin_angle2 + y * cos_angle2, z
new_orientation = q1 * q2
target = Box(center=new_center, size=size, orientation=new_orientation)
assert original.rotate_around_origin(q2) == target
def transform_box_from_world_to_flat_vehicle_coordinates(first_train_sample_box: Box, sample_data_token: str,
lyftd: LyftDataset):
sample_box = first_train_sample_box.copy()
sd_record = lyftd.get("sample_data", sample_data_token)
cs_record = lyftd.get("calibrated_sensor", sd_record["calibrated_sensor_token"])
sensor_record = lyftd.get("sensor", cs_record["sensor_token"])
pose_record = lyftd.get("ego_pose", sd_record["ego_pose_token"])
# Move box to ego vehicle coord system parallel to world z plane
ypr = Quaternion(pose_record["rotation"]).yaw_pitch_roll
yaw = ypr[0]
sample_box.translate(-np.array(pose_record["translation"]))
sample_box.rotate(Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)]).inverse)
return sample_box
def get_box_corners(transformed_box: Box,
cam_intrinsic_mtx: np.array,
frustum_pointnet_convention=True):
box_corners_on_cam_coord = transformed_box.corners()
# Regarrange to conform Frustum-pointnet's convention
if frustum_pointnet_convention:
rearranged_idx = [0, 3, 7, 4, 1, 2, 6, 5]
box_corners_on_cam_coord = box_corners_on_cam_coord[:, rearranged_idx]
assert np.allclose((box_corners_on_cam_coord[:, 0] + box_corners_on_cam_coord[:, 6]) / 2,
np.array(transformed_box.center))
# For perspective transformation, the normalization should set to be True
box_corners_on_image = view_points(box_corners_on_cam_coord, view=cam_intrinsic_mtx, normalize=True)
return box_corners_on_image
def _info_to_nusc_box(info):
from lyft_dataset_sdk.utils.data_classes import Box
box3d = info['gt_boxes'].copy()
names = info['gt_names'].copy()
box3d[:, 6] = -box3d[:, 6] - np.pi / 2
box_list = []
for i in range(box3d.shape[0]):
quat = Quaternion(axis=[0, 0, 1], radians=box3d[i, 6])
velocity = (np.nan, np.nan, np.nan)
box = Box(box3d[i, :3],
box3d[i, 3:6],
quat,
name=names[i],
velocity=velocity)
box_list.append(box)
return box_list
def parse_string_to_box(ps, with_score=True, output_type="box", sample_token=None) -> List[Box]:
boxes = []
col_num = 8
if with_score:
col_num = 9
object_params = ps.split()
n_objects = len(object_params)
for i in range(n_objects // col_num):
if with_score:
score, x, y, z, w, l, h, yaw, c = tuple(object_params[i * 9: (i + 1) * 9])
score = float(score)
else:
x, y, z, w, l, h, yaw, c = tuple(object_params[i * 8: (i + 1) * 8])
score = 1.0 # assume ground truth
if not (float(w) > 0 and float(l) > 0 and float(h) > 0):
warnings.warn("wrong wlh value")
continue
orient_q = Quaternion(axis=[0, 0, 1], angle=float(yaw))
center_pos = [float(x), float(y), float(z)]
wlh = [float(w), float(l), float(h)]
obj_name = c
if output_type == "3dbox":
boxes.append(Box3D(translation=center_pos, size=wlh, rotation=orient_q.q, name=obj_name, score=score,
sample_token=sample_token))
elif output_type == "box":
boxes.append(Box(center=center_pos, size=wlh, orientation=orient_q, name=obj_name, score=score))
elif output_type == "dict":
boxes.append({'translation': center_pos, 'size': wlh, 'rotation': orient_q.q,
'name': obj_name, 'score': score, 'sample_token': sample_token})
else:
raise ValueError("output_type must be either 3dbox, box, or dict")
return boxes
def convert_boxes3d_xyzlwhr_to_Box(boxes3d_xyzlwhr,
token=None,
name=None,
score=0.0):
x = boxes3d_xyzlwhr[:, 0]
y = boxes3d_xyzlwhr[:, 1]
z = boxes3d_xyzlwhr[:, 2]
l = boxes3d_xyzlwhr[:, 3]
w = boxes3d_xyzlwhr[:, 4]
h = boxes3d_xyzlwhr[:, 5]
rz = boxes3d_xyzlwhr[:, 6]
M_rz = np.array([
[np.cos(rz), -np.sin(rz), np.zeros_like(rz)],
[np.sin(rz), np.cos(rz), np.zeros_like(rz)],
[np.zeros_like(rz),
np.zeros_like(rz),
np.ones_like(rz)],
])
M_rz = M_rz.transpose(2, 0, 1)
M_rz = R.from_dcm(M_rz)
quat = M_rz.as_quat()
# XYZW -> WXYZ order of elements
quat = quat[:, [3, 0, 1, 2]]
boxes3d = []
for i in range(len(boxes3d_xyzlwhr)):
box3d = Box(token=token,
center=[x[i], y[i], z[i]],
size=[w[i], l[i], h[i]],
orientation=Quaternion(quat[i]),
name=name,
score=score)
boxes3d.append(box3d)
return boxes3d
def render_sample_data(sample_data_token: str,
with_anns: bool = True,
box_vis_level: BoxVisibility = BoxVisibility.ANY,
axes_limit: float = 40,
ax: Axes = None,
num_sweeps: int = 1,
out_path: str = None,
underlay_map: bool = False,
detections: list = [],
categories: list = [],
valLyft: LyftDataset = None):
"""Render sample data onto axis.
Args:
sample_data_token: Sample_data token.
with_anns: Whether to draw annotations.
box_vis_level: If sample_data is an image, this sets required visibility for boxes.
axes_limit: Axes limit for lidar and radar (measured in meters).
ax: Axes onto which to render.
num_sweeps: Number of sweeps for lidar and radar.
out_path: Optional path to save the rendered figure to disk.
underlay_map: When set to true, LIDAR data is plotted onto the map. This can be slow.
"""
# Get sensor modality.
sd_record = valLyft.get("sample_data", sample_data_token)
sensor_modality = sd_record["sensor_modality"]
if sensor_modality == "camera":
# Load boxes and image.
data_path, _, camera_intrinsic = valLyft.get_sample_data(
sample_data_token, box_vis_level=box_vis_level)
data = Image.open(data_path)
# Init axes.
if ax is None:
_, ax = plt.subplots(1, 1, figsize=(9, 16))
# Show image.
ax.imshow(data)
#categories = ['car', 'pedestrian', 'truck', 'bicycle', 'bus', 'other_vehicle', 'motorcycle', 'emergency_vehicle', 'animal']
# Show boxes.
if with_anns:
boxes = []
for c1, detection in enumerate(detections):
#print(categories)
cat = categories[c1]
#print(cat)
#import pdb; pdb.set_trace()
box = Box(detection[:3],
detection[3:6],
Quaternion(np.array(detection[6:10])),
name=cat)
boxes.append(box)
for box in boxes:
c = np.array(get_color(box.name)) / 255.0
box.render(ax,
view=camera_intrinsic,
normalize=True,
colors=(c, c, c))
# Limit visible range.
ax.set_xlim(0, data.size[0])
ax.set_ylim(data.size[1], 0)
ax.axis("off")
ax.set_title(sd_record["channel"])
ax.set_aspect("equal")
if out_path is not None:
num = len([name for name in os.listdir(out_path)])
out_path = out_path + str(num).zfill(
5) + "_" + sample_data_token + ".png"
plt.savefig(out_path)
plt.close("all")
return out_path
def main():
tic = time.time()
if params.debug:
logger = init_logger('_log/03_make_sub_debug.log', level=10)
else:
logger = init_logger('_log/03_make_sub.log', level=20)
level5data = LyftDataset(data_path='../../dataset/test',
json_path='../../dataset/test/data/',
verbose=True)
sub_pre = pd.read_csv(params.path_sub_pre)
print(sub_pre.head())
sample = pd.read_csv('../../dataset/sample_submission.csv')
print(sample.head())
target_tokens = sample['Id']
if params.debug:
target_tokens = target_tokens[:20]
list_subs = list()
img_size = 2048
lidar_range = 100
for i, target_token in enumerate(tqdm(target_tokens)):
target_subs = sub_pre.query('token==@target_token')
list_sub_token = list()
for _, target_sub in target_subs.iterrows():
x = target_sub['x']
y = target_sub['y']
z = target_sub['z']
length = target_sub['length']
width = target_sub['width']
height = target_sub['height']
rotate = target_sub['rotate']
width = width * (lidar_range * 2) / (img_size - 1)
length = length * (lidar_range * 2) / (img_size - 1)
height = height * (lidar_range * 2) / (img_size - 1)
x = x * (lidar_range * 2) / (img_size - 1) - lidar_range
y = y * (lidar_range * 2) / (img_size - 1) - lidar_range
z = z * (lidar_range * 2) / (img_size - 1) - lidar_range
rotate = -rotate
quat = Quaternion(math.cos(rotate / 2), 0, 0, math.sin(rotate / 2))
print(quat.yaw_pitch_roll)
pred_box = Box([x, y, z], [width, length, height], quat)
my_sample = level5data.get('sample', target_token)
rev_token = level5data.get('sample_data',
my_sample['data']['LIDAR_TOP'])['token']
rev_box = reverse_box(pred_box, level5data, rev_token)
sub_i = '{:.9f} '.format(target_sub['confidence']) + \
' '.join(['{:.3f}'.format(v) for v in rev_box.center]) + \
' ' + ' '.join(['{:.3f}'.format(v) for v in rev_box.wlh]) + \
' {:.3f}'.format(rev_box.orientation.yaw_pitch_roll[0]) + ' {}'.format(target_sub['name'])
logger.debug('sub_i')
logger.debug(sub_i)
list_sub_token.append(sub_i)
if len(list_sub_token) == 0:
sub_token = ''
else:
sub_token = ' '.join(list_sub_token)
logger.info('submit token !')
logger.info(sub_token)
list_subs.append(sub_token)
submission = pd.DataFrame()
submission['Id'] = target_tokens
submission['PredictionString'] = list_subs
dir_sub = Path('_submission')
dir_sub.mkdir(exist_ok=True)
submission.to_csv(dir_sub / params.path_sub_pre.name, index=False)
logger.info('elapsed time: {:.1f} [min]'.format(
(time.time() - tic) / 60.0))
def get_boxes(self,
token: str,
filter_classes: List[str] = None,
max_dist: float = None) -> List[Box]:
"""Load up all the boxes associated with a sample.
Boxes are in nuScenes lidar frame.
Args:
token: KittiDB unique id.
filter_classes: List of Kitti classes to use or None to use all.
max_dist: List of Kitti classes to use or None to use all.
Returns: Boxes in nuScenes lidar reference frame.
"""
# Get transforms for this sample
transforms = self.get_transforms(token, root=self.root)
boxes = []
if token.startswith("test_"):
# No boxes to return for the test set.
return boxes
with open(KittiDB.get_filepath(token, "label_2", root=self.root),
"r") as f:
for line in f:
# Parse this line into box information.
parsed_line = self.parse_label_line(line)
if parsed_line["name"] in {"DontCare", "Misc"}:
continue
center = parsed_line["xyz_camera"]
wlh = parsed_line["wlh"]
yaw_camera = parsed_line["yaw_camera"]
name = parsed_line["name"]
score = parsed_line["score"]
# Optional: Filter classes.
if filter_classes is not None and name not in filter_classes:
continue
# The Box class coord system is oriented the same way as as KITTI LIDAR: x forward, y left, z up.
# For orientation confer: http://www.cvlibs.net/datasets/kitti/setup.php.
# 1: Create box in Box coordinate system with center at origin.
# The second quaternion in yaw_box transforms the coordinate frame from the object frame
# to KITTI camera frame. The equivalent cannot be naively done afterwards, as it's a rotation
# around the local object coordinate frame, rather than the camera frame.
quat_box = Quaternion(axis=(0, 1, 0),
angle=yaw_camera) * Quaternion(
axis=(1, 0, 0), angle=np.pi / 2)
box = Box([0.0, 0.0, 0.0], wlh, quat_box, name=name)
# 2: Translate: KITTI defines the box center as the bottom center of the vehicle. We use true center,
# so we need to add half height in negative y direction, (since y points downwards), to adjust. The
# center is already given in camera coord system.
box.translate(center + np.array([0, -wlh[2] / 2, 0]))
# 3: Transform to KITTI LIDAR coord system. First transform from rectified camera to camera, then
# camera to KITTI lidar.
box.rotate(Quaternion(matrix=transforms["r0_rect"]).inverse)
box.translate(-transforms["velo_to_cam"]["T"])
box.rotate(
Quaternion(matrix=transforms["velo_to_cam"]["R"]).inverse)
# 4: Transform to nuScenes LIDAR coord system.
box.rotate(self.kitti_to_nu_lidar)
# Set score or NaN.
box.score = score
# Set dummy velocity.
box.velocity = np.array((0.0, 0.0, 0.0))
# Optional: Filter by max_dist
if max_dist is not None:
dist = np.sqrt(np.sum(box.center[:2]**2))
if dist > max_dist:
continue
boxes.append(box)
return boxes
def box_nuscenes_to_kitti(
box: Box,
velo_to_cam_rot: Quaternion,
velo_to_cam_trans: np.ndarray,
r0_rect: Quaternion,
kitti_to_nu_lidar_inv: Quaternion = Quaternion(axis=(0, 0, 1),
angle=np.pi /
2).inverse,
) -> Box:
"""Transform from nuScenes lidar frame to KITTI reference frame.
Args:
box: Instance in nuScenes lidar frame.
velo_to_cam_rot: Quaternion to rotate from lidar to camera frame.
velo_to_cam_trans: <np.float: 3>. Translate from lidar to camera frame.
r0_rect: Quaternion to rectify camera frame.
kitti_to_nu_lidar_inv: Quaternion to rotate nuScenes to KITTI LIDAR.
Returns: Box instance in KITTI reference frame.
"""
# Copy box to avoid side-effects.
box = box.copy()
# Rotate to KITTI lidar.
box.rotate(kitti_to_nu_lidar_inv)
# Transform to KITTI camera.
box.rotate(velo_to_cam_rot)
box.translate(velo_to_cam_trans)
# Rotate to KITTI rectified camera.
box.rotate(r0_rect)
# KITTI defines the box center as the bottom center of the object.
# We use the true center, so we need to adjust half height in y direction.
box.translate(np.array([0, box.wlh[2] / 2, 0]))
return box
def get_pred_dict(pred, sample_token, classes, token2info):
# only for nuscenes
attribute_NameMapping = {
'bicycle': 'cycle.with_rider',
'bus': 'vehicle.moving',
'car': 'vehicle.moving',
'other_vehicle': 'vehicle.moving',
'construction_vehicle': 'vehicle.moving',
'motorcycle': 'cycle.with_rider',
'pedestrian': 'pedestrian.moving',
'trailer': 'vehicle.moving',
'truck': 'vehicle.moving'
}
detection_NameMapping = {
'bicycle': 'bicycle',
'bus': 'bus',
'car': 'car',
'other_vehicle': 'car',
'construction_vehicle': 'construction_vehicle',
'motorcycle': 'motorcycle',
'pedestrian': 'pedestrian',
'trailer': 'trailer',
'truck': 'truck',
'barrier': 'barrier',
'traffic_cone': 'traffic_cone'
}
boxes_lidar = pred["box3d_lidar"]
boxes_class = pred["label_preds"]
scores = pred['scores']
preds_classes = [classes[x] for x in boxes_class]
box_centers = boxes_lidar[:, :3]
box_yaws = boxes_lidar[:, -1]
box_wlh = boxes_lidar[:, 3:6]
info = token2info[sample_token] # a `sample` token
boxes = []
pred_str = ''
dict_frame_detections = {str(sample_token): []}
for idx in range(len(boxes_lidar)):
translation = box_centers[idx]
yaw = -box_yaws[idx] - pi / 2
size = box_wlh[idx]
name = preds_classes[idx]
detection_score = scores[idx]
quat = Quaternion(scalar=np.cos(yaw / 2),
vector=[0, 0, np.sin(yaw / 2)])
box = Box(center=box_centers[idx],
size=size,
orientation=quat,
score=detection_score,
name=name,
token=sample_token)
box = to_glb(box, info)
x = box.center[0]
y = box.center[1]
z = box.center[2]
w = np.float64(box.wlh[0])
l = np.float64(box.wlh[1])
h = np.float64(box.wlh[2])
score = np.float64(box.score)
q = Quaternion(axis=[0, 0, 1], angle=box.orientation.yaw_pitch_roll[0])
qw = q.w
qx = q.x
qy = q.y
qz = q.z
#print('--')
#print(f'{type(x)} {type(y)} {type(z)}')
#print(f'{type(score)}')
#print(f'{type(w)} {type(l)} {type(h)}')
#print(f'{type(qw)} {type(qx)} {type(qy)} {type(qz)}')
#print(type(sample_token))
dict_frame_detection = {
'sample_token': sample_token,
'translation': [x, y, z],
'size': [w, l, h],
'rotation': [qw, qx, qy, qz],
'velocity': [0, 0],
'detection_name': detection_NameMapping[name],
'detection_score': score,
'attribute_name': attribute_NameMapping[name]
}
dict_frame_detections[str(sample_token)].append(dict_frame_detection)
return dict_frame_detections
