def get_heading_loss(self, heading_scores, heading_res_norm, heading_class_label, heading_res_norm_label):
heading_class_loss = F.cross_entropy(heading_scores, heading_class_label)
# b, NUM_HEADING_BIN -> b, 1
heading_res_norm_select = torch.gather(heading_res_norm, 1, heading_class_label.view(-1, 1))
heading_res_norm_loss = huber_loss(
heading_res_norm_select.squeeze(1) - heading_res_norm_label, delta=1.0)
return heading_class_loss, heading_res_norm_loss
def get_size_loss(self, size_scores, size_res_norm, size_class_label, size_res_label_norm):
batch_size = size_scores.shape[0]
size_class_loss = F.cross_entropy(size_scores, size_class_label)
# b, NUM_SIZE_CLUSTER, 3 -> b, 1, 3
size_res_norm_select = torch.gather(size_res_norm, 1,
size_class_label.view(batch_size, 1, 1).expand(
batch_size, 1, 3))
size_norm_dist = torch.norm(
size_res_label_norm - size_res_norm_select.squeeze(1), 2, dim=-1)
size_res_norm_loss = huber_loss(size_norm_dist, delta=1.0)
return size_class_loss, size_res_norm_loss
def get_corner_loss(self, preds, gts):
center_label, heading_label, size_label = gts
center_preds, heading_preds, size_preds = preds
corners_3d_gt = get_box3d_corners_helper(center_label, heading_label, size_label)
corners_3d_gt_flip = get_box3d_corners_helper(center_label, heading_label + np.pi, size_label)
corners_3d_pred = get_box3d_corners_helper(center_preds, heading_preds, size_preds)
# N, 8, 3
corners_dist = torch.min(
torch.norm(corners_3d_pred - corners_3d_gt, 2, dim=-1).mean(-1),
torch.norm(corners_3d_pred - corners_3d_gt_flip, 2, dim=-1).mean(-1))
# corners_dist = torch.norm(corners_3d_pred - corners_3d_gt, 2, dim=-1)
corners_loss = huber_loss(corners_dist, delta=1.0)
return corners_loss, corners_3d_gt
def get_center_loss(self, pred_offsets, gt_offsets):
center_dist = torch.norm(gt_offsets - pred_offsets, 2, dim=-1)
center_loss = huber_loss(center_dist, delta=3.0)
return center_loss
