def assign_targets(self, batch_dict):
batch_size = batch_dict['batch_size']
with torch.no_grad():
targets_dict = self.proposal_target_layer.forward(batch_dict)
rois = targets_dict['rois'] # (B, N, 7 + C)
gt_of_rois = targets_dict['gt_of_rois'] # (B, N, 7 + C + 1)
targets_dict['gt_of_rois_src'] = gt_of_rois.clone().detach()
# canonical transformation
roi_center = rois[:, :, 0:3]
roi_ry = rois[:, :, 6] % (2 * np.pi)
gt_of_rois[:, :, 0:3] = gt_of_rois[:, :, 0:3] - roi_center
gt_of_rois[:, :, 6] = gt_of_rois[:, :, 6] - roi_ry
# transfer LiDAR coords to local coords
gt_of_rois = common_utils.rotate_points_along_z(
points=gt_of_rois.view(-1, 1, gt_of_rois.shape[-1]), angle=-roi_ry.view(-1)
).view(batch_size, -1, gt_of_rois.shape[-1])
# flip orientation if rois have opposite orientation
heading_label = gt_of_rois[:, :, 6] % (2 * np.pi) # 0 ~ 2pi
opposite_flag = (heading_label > np.pi * 0.5) & (heading_label < np.pi * 1.5)
heading_label[opposite_flag] = (heading_label[opposite_flag] + np.pi) % (2 * np.pi) # (0 ~ pi/2, 3pi/2 ~ 2pi)
flag = heading_label > np.pi
heading_label[flag] = heading_label[flag] - np.pi * 2 # (-pi/2, pi/2)
heading_label = torch.clamp(heading_label, min=-np.pi / 2, max=np.pi / 2)
gt_of_rois[:, :, 6] = heading_label
targets_dict['gt_of_rois'] = gt_of_rois
return targets_dict
def boxes_to_corners_3d(boxes3d):
"""
7 -------- 4
/| /|
6 -------- 5 .
| | | |
. 3 -------- 0
|/ |/
2 -------- 1
Args:
boxes3d: (N, 7) [x, y, z, dx, dy, dz, heading], (x, y, z) is the box center
Returns:
"""
boxes3d, is_numpy = common_utils.check_numpy_to_torch(boxes3d)
template = boxes3d.new_tensor((
[1, 1, -1],
[1, -1, -1],
[-1, -1, -1],
[-1, 1, -1],
[1, 1, 1],
[1, -1, 1],
[-1, -1, 1],
[-1, 1, 1],
)) / 2
corners3d = boxes3d[:, None, 3:6].repeat(1, 8, 1) * template[None, :, :]
corners3d = common_utils.rotate_points_along_z(corners3d.view(-1, 8, 3),
boxes3d[:,
6]).view(-1, 8, 3)
corners3d += boxes3d[:, None, 0:3]
return corners3d.numpy() if is_numpy else corners3d
def generate_predicted_boxes(self, batch_size, rois, cls_preds, box_preds):
"""
Args:
batch_size:
rois: (B, N, 7)
cls_preds: (BN, num_class)
box_preds: (BN, code_size)
Returns:
"""
code_size = self.box_coder.code_size
# batch_cls_preds: (B, N, num_class or 1)
batch_cls_preds = cls_preds.view(batch_size, -1, cls_preds.shape[-1])
batch_box_preds = box_preds.view(batch_size, -1, code_size)
roi_ry = rois[:, :, 6].view(-1)
roi_xyz = rois[:, :, 0:3].view(-1, 3)
local_rois = rois.clone().detach()
local_rois[:, :, 0:3] = 0
batch_box_preds = self.box_coder.decode_torch(batch_box_preds, local_rois).view(-1, code_size)
batch_box_preds = common_utils.rotate_points_along_z(
batch_box_preds.unsqueeze(dim=1), roi_ry
).squeeze(dim=1)
batch_box_preds[:, 0:3] += roi_xyz
batch_box_preds = batch_box_preds.view(batch_size, -1, code_size)
return batch_cls_preds, batch_box_preds
def get_global_grid_points_of_roi(self, rois, grid_size):
rois = rois.view(-1, rois.shape[-1])
batch_size_rcnn = rois.shape[0]
local_roi_grid_points = self.get_dense_grid_points(rois, batch_size_rcnn, grid_size) # (B, 6x6x6, 3)
global_roi_grid_points = common_utils.rotate_points_along_z(
local_roi_grid_points.clone(), rois[:, 6]
).squeeze(dim=1)
global_center = rois[:, 0:3].clone()
global_roi_grid_points += global_center.unsqueeze(dim=1)
return global_roi_grid_points, local_roi_grid_points
def generate_predict_box(self,batch_size,cls_preds,reg_preds,rois):
cls_preds = cls_preds.view(batch_size,-1,cls_preds.shape[-1])
reg_preds = reg_preds.view(batch_size,-1,self.code_func.code_size)
rois_y = rois[...,6].view(-1)
rois_xyz = rois[...,0:3]
rois_temp = rois.clone().detach()
rois_temp[...,0:3] = 0
box_ct = self.code_func.decode_torch(reg_preds,rois_temp).view(-1,self.code_func.code_size)
box_preds = common_utils.rotate_points_along_z(
box_ct.unsqueeze(dim=1),rois_y).squeeze(dim=1).view(batch_size,-1,self.code_func.code_size)
box_preds[...,0:3] += rois_xyz
return cls_preds,box_preds
def get_reg_loss(self):
loss_cfgs = self.model_cfg.LOSS_CONFIG
reg_preds = self.forward_ret_dict["rcnn_reg"]
reg_targets = self.forward_ret_dict['rcnn_reg_targets']
reg_valid_mask = self.forward_ret_dict['fg_valid_reg_mask']
batch_size = reg_targets.shape[0]
reg_preds = reg_preds.view(batch_size, -1, reg_preds.shape[-1])
fg_mask = (reg_valid_mask>0)
try:
fg_num = fg_mask.cpu().long().sum().item()
except:
raise ValueError
tb_dict = {"rcnn_fg_reg_num":fg_num}
loss = 0.0
if loss_cfgs.REG_LOSS == "smooth-l1":
reg_loss = self.reg_loss_func(reg_preds,reg_targets)
fg_flag = (fg_mask>0).float().unsqueeze(dim=-1)
reg_loss = (reg_loss*fg_flag).sum()/torch.clamp(fg_flag.sum(),1.0)
reg_loss = reg_loss*loss_cfgs.LOSS_WEIGHTS["rcnn_reg_weight"]
tb_dict["rcnn_fg_reg_loss"] = reg_loss.item()
loss +=reg_loss
if loss_cfgs.CORNER_LOSS_REGULARIZATION and fg_num>0:
# reg_valid_mask[...,0:3] =1.0 for debug
corner_targets = self.forward_ret_dict['rcnn_reg_corners_target']
fg_reg_preds = reg_preds[fg_mask]
rois = self.forward_ret_dict["rois"]
fg_rois = rois[fg_mask]
fg_rois_ry = fg_rois[...,6]
fg_rois_xyz = fg_rois[...,0:3]
fg_anchors = fg_rois.clone().detach()
fg_anchors[...,0:3] = 0
fg_box_preds_ct = self.code_func.decode_torch(fg_reg_preds,fg_anchors)
fg_box_preds = common_utils.rotate_points_along_z(fg_box_preds_ct.unsqueeze(dim=1),fg_rois_ry).squeeze(dim=1)
fg_box_preds[...,0:3] += fg_rois_xyz
fg_corner_tagrets = corner_targets[fg_mask]
corner_loss = self.corner_loss_func(fg_box_preds[...,0:7],fg_corner_tagrets[...,0:7])
corner_loss = corner_loss.mean()
corner_loss = corner_loss*loss_cfgs.LOSS_WEIGHTS["rcnn_corner_weight"]
loss += corner_loss
tb_dict.update({"rcnn_corner_loss":corner_loss.item()})
return loss,tb_dict
def roi_grid_pool_layer(self,batch_dict):
batch_size = batch_dict["rois"].shape[0]
rois = batch_dict["rois"].reshape(-1,batch_dict["rois"].shape[-1])
point_coords = batch_dict["point_coords"]
point_features = batch_dict['point_features']
point_features = point_features* batch_dict["point_cls_scores"].view(-1,1)
grid_size = self.model_cfg.ROI_GRID_POOL.GRID_SIZE
#生成rois的网格点将每个roi分成6*6*6的网格
rcnn_batch_size = rois.shape[0]
local_grid = rois.new_ones(grid_size,grid_size,grid_size)
local_grid_id = torch.nonzero(local_grid).float()
local_grid_id = local_grid_id.repeat(rcnn_batch_size,1,1)
rois_size = rois[...,3:6].unsqueeze(dim=1).clone()
gloal_grid_points = (local_grid_id+0.5)*(rois_size/grid_size) - rois_size/2
rois_ry = rois[...,6].clone()
gloal_grid_points = common_utils.rotate_points_along_z(
gloal_grid_points.clone(),rois_ry.view(-1))
roi_centers = rois[...,0:3].clone()
gloal_grid_points += roi_centers.unsqueeze(dim=1)
gloal_grid_points = gloal_grid_points.view(batch_size,-1,gloal_grid_points.shape[-1])
xyz = point_coords[:,1:4]
xyz_batch_count = xyz.new_zeros(batch_size,).int()
for bn_idx in range(batch_size):
xyz_batch_count[bn_idx] = (point_coords[:,0] == bn_idx).sum()
new_xyz = gloal_grid_points.view(-1,3)
new_xyz_bn_count = new_xyz.new_zeros(batch_size).fill_(gloal_grid_points.shape[1]).int()
pooled_xyz,pooled_features = self.roi_pool_layer(
xyz.contiguous(),
xyz_batch_count,
new_xyz,
new_xyz_bn_count,
features = point_features
)
return pooled_features
def assign_stack_targets(self,
points,
gt_boxes,
extend_gt_boxes=None,
ret_box_labels=False,
ret_part_labels=False,
set_ignore_flag=True,
use_ball_constraint=False,
central_radius=2.0):
"""
Args:
points: (N1 + N2 + N3 + ..., 4) [bs_idx, x, y, z]
gt_boxes: (B, M, 8)
extend_gt_boxes: [B, M, 8]
ret_box_labels:
ret_part_labels:
set_ignore_flag:
use_ball_constraint:
central_radius:
Returns:
point_cls_labels: (N1 + N2 + N3 + ...), long type, 0:background, -1:ignored
point_box_labels: (N1 + N2 + N3 + ..., code_size)
"""
assert len(
points.shape
) == 2 and points.shape[1] == 4, 'points.shape=%s' % str(points.shape)
assert len(gt_boxes.shape) == 3 and gt_boxes.shape[
2] == 8, 'gt_boxes.shape=%s' % str(gt_boxes.shape)
assert extend_gt_boxes is None or len(extend_gt_boxes.shape) == 3 and extend_gt_boxes.shape[2] == 8, \
'extend_gt_boxes.shape=%s' % str(extend_gt_boxes.shape)
assert set_ignore_flag != use_ball_constraint, 'Choose one only!'
batch_size = gt_boxes.shape[0]
bs_idx = points[:, 0]
point_cls_labels = points.new_zeros(points.shape[0]).long()
point_box_labels = gt_boxes.new_zeros(
(points.shape[0], 8)) if ret_box_labels else None
point_part_labels = gt_boxes.new_zeros(
(points.shape[0], 3)) if ret_part_labels else None
for k in range(batch_size):
bs_mask = (bs_idx == k)
points_single = points[bs_mask][:, 1:4]
point_cls_labels_single = point_cls_labels.new_zeros(bs_mask.sum())
box_idxs_of_pts = roiaware_pool3d_utils.points_in_boxes_gpu(
points_single.unsqueeze(dim=0),
gt_boxes[k:k + 1, :, 0:7].contiguous()).long().squeeze(dim=0)
box_fg_flag = (box_idxs_of_pts >= 0)
if set_ignore_flag:
extend_box_idxs_of_pts = roiaware_pool3d_utils.points_in_boxes_gpu(
points_single.unsqueeze(dim=0),
extend_gt_boxes[k:k + 1, :,
0:7].contiguous()).long().squeeze(dim=0)
fg_flag = box_fg_flag
ignore_flag = fg_flag ^ (extend_box_idxs_of_pts >= 0)
point_cls_labels_single[ignore_flag] = -1
elif use_ball_constraint:
box_centers = gt_boxes[k][box_idxs_of_pts][:, 0:3].clone()
box_centers[:, 2] += gt_boxes[k][box_idxs_of_pts][:, 5] / 2
ball_flag = ((box_centers - points_single).norm(dim=1) <
central_radius)
fg_flag = box_fg_flag & ball_flag
else:
raise NotImplementedError
gt_box_of_fg_points = gt_boxes[k][box_idxs_of_pts[fg_flag]]
point_cls_labels_single[
fg_flag] = 1 if self.num_class == 1 else gt_box_of_fg_points[:,
7].long(
)
point_cls_labels[bs_mask] = point_cls_labels_single
if ret_box_labels:
point_box_labels_single = point_box_labels.new_zeros(
(bs_mask.sum(), 8))
fg_point_box_labels = self.box_coder.encode_torch(
points_single[fg_flag], gt_box_of_fg_points)
point_box_labels_single[fg_flag] = fg_point_box_labels
point_box_labels[bs_mask] = point_box_labels_single
if ret_part_labels:
point_part_labels_single = point_part_labels.new_zeros(
(bs_mask.sum(), 3))
transformed_points = points_single[
fg_flag] - gt_box_of_fg_points[:, 0:3]
transformed_points = common_utils.rotate_points_along_z(
transformed_points.view(-1, 1, 3),
-gt_box_of_fg_points[:, 6]).view(-1, 3)
offset = torch.tensor([0.5, 0.5, 0.5
]).view(1,
3).type_as(transformed_points)
point_part_labels_single[fg_flag] = (
transformed_points / gt_box_of_fg_points[:, 3:6]) + offset
point_part_labels[bs_mask] = point_part_labels_single
targets_dict = {
'point_cls_labels': point_cls_labels,
'point_box_labels': point_box_labels,
'point_part_labels': point_part_labels
}
return targets_dict
def forward(self, batch_dict):
batch_rois_src,batch_roi_labels_src,batch_rois_scores_src,\
batch_gt_of_rois_src,batch_rois_max_iou3d_src = self.get_sampled_rois(batch_dict)
#生成回归的mask
fg_valid_reg_mask = (batch_rois_max_iou3d_src >
self.target_cfg.REG_FG_THRESH).long()
#生成iou为引导的roi, score作为rcnn阶段类别损失函数的target
rcnn_cls_targets = batch_rois_max_iou3d_src.clone().detach()
batch_size = batch_dict["batch_size"]
fg_mask = rcnn_cls_targets > self.target_cfg.CLS_FG_THRESH
bg_mask = rcnn_cls_targets < self.target_cfg.CLS_BG_THRESH
interval_mask = (fg_mask == 0) & (bg_mask == 0)
rcnn_cls_targets[fg_mask] = 1.0
rcnn_cls_targets[bg_mask] = 0.0
rcnn_cls_targets[interval_mask] = (rcnn_cls_targets[interval_mask] -
0.25) * 2
rcnn_cls_targets = rcnn_cls_targets.float()
#局部坐标变换
batch_gt_of_rois = batch_gt_of_rois_src.clone().detach()
batch_rois = batch_rois_src.clone().detach()
roi_centers = batch_rois[..., 0:3]
roi_rot = batch_rois[..., 6]
batch_gt_of_rois[..., 6] = batch_gt_of_rois[..., 6] % (2 * np.pi)
batch_gt_of_rois[..., 6] = batch_gt_of_rois[..., 6] - roi_rot
batch_gt_of_rois[..., 0:3] = batch_gt_of_rois[..., 0:3] - roi_centers
batch_gt_of_rois = common_utils.rotate_points_along_z(
batch_gt_of_rois.reshape(-1, 1, batch_gt_of_rois.shape[-1]),
-roi_rot.view(-1))
#对gt_of_roi的角度进行变换,
gt_of_roi_ry = batch_gt_of_rois[..., 6] % (np.pi * 2) #(0-2*pi)
opposite_flag = (gt_of_roi_ry > np.pi * 0.5) & (gt_of_roi_ry <
np.pi * 1.5)
gt_of_roi_ry[opposite_flag] = (gt_of_roi_ry[opposite_flag] - np.pi) % (
np.pi * 2) #(0 ~ pi*0.5, 1.5*pi ~ 2*pi)
flag = gt_of_roi_ry > np.pi
gt_of_roi_ry[flag] = gt_of_roi_ry[flag] - np.pi * 2 #(-pi*0.5,pi*0.5)
gt_of_roi_ry = torch.clamp(gt_of_roi_ry, min=-np.pi / 2, max=np.pi / 2)
batch_gt_of_rois[..., 6] = gt_of_roi_ry
batch_gt_of_rois = batch_gt_of_rois.reshape(
batch_size, -1, self.code_func.code_size + 1)
#生成回归目标
batch_rois[..., 0:3] = 0
batch_rois[..., 6] = 0
rcnn_reg_targets = self.code_func.encode_torch(
batch_gt_of_rois[..., :7], batch_rois)
#生成corner loss的targets
rcnn_reg_corners_target = batch_gt_of_rois_src
ret_dict = {
"rois": batch_rois_src,
"roi_scores": batch_rois_scores_src,
"roi_labels": batch_roi_labels_src,
"gt_of_rois_src": batch_gt_of_rois_src,
"gt_of_rois": batch_gt_of_rois,
"roi_iou3d": batch_rois_max_iou3d_src,
"fg_valid_reg_mask": fg_valid_reg_mask,
"rcnn_reg_targets": rcnn_reg_targets,
"rcnn_reg_corners_target": rcnn_reg_corners_target,
"rcnn_cls_targets": rcnn_cls_targets
}
return ret_dict
def get_box_reg_layer_loss(self, forward_ret_dict):
loss_cfgs = self.model_cfg.LOSS_CONFIG
code_size = self.box_coder.code_size
reg_valid_mask = forward_ret_dict['reg_valid_mask'].view(-1)
gt_boxes3d_ct = forward_ret_dict['gt_of_rois'][..., 0:code_size]
gt_of_rois_src = forward_ret_dict['gt_of_rois_src'][..., 0:code_size].view(-1, code_size)
rcnn_reg = forward_ret_dict['rcnn_reg'] # (rcnn_batch_size, C)
roi_boxes3d = forward_ret_dict['rois']
rcnn_batch_size = gt_boxes3d_ct.view(-1, code_size).shape[0]
fg_mask = (reg_valid_mask > 0)
fg_sum = fg_mask.long().sum().item()
tb_dict = {}
if loss_cfgs.REG_LOSS == 'smooth-l1':
rois_anchor = roi_boxes3d.clone().detach().view(-1, code_size)
rois_anchor[:, 0:3] = 0
rois_anchor[:, 6] = 0
reg_targets = self.box_coder.encode_torch(
gt_boxes3d_ct.view(rcnn_batch_size, code_size), rois_anchor
)
rcnn_loss_reg = self.reg_loss_func(
rcnn_reg.view(rcnn_batch_size, -1).unsqueeze(dim=0),
reg_targets.unsqueeze(dim=0),
) # [B, M, 7]
rcnn_loss_reg = (rcnn_loss_reg.view(rcnn_batch_size, -1) * fg_mask.unsqueeze(dim=-1).float()).sum() / max(fg_sum, 1)
rcnn_loss_reg = rcnn_loss_reg * loss_cfgs.LOSS_WEIGHTS['rcnn_reg_weight']
tb_dict['rcnn_loss_reg'] = rcnn_loss_reg.item()
if loss_cfgs.CORNER_LOSS_REGULARIZATION and fg_sum > 0:
# TODO: NEED to BE CHECK
fg_rcnn_reg = rcnn_reg.view(rcnn_batch_size, -1)[fg_mask]
fg_roi_boxes3d = roi_boxes3d.view(-1, code_size)[fg_mask]
fg_roi_boxes3d = fg_roi_boxes3d.view(1, -1, code_size)
batch_anchors = fg_roi_boxes3d.clone().detach()
roi_ry = fg_roi_boxes3d[:, :, 6].view(-1)
roi_xyz = fg_roi_boxes3d[:, :, 0:3].view(-1, 3)
batch_anchors[:, :, 0:3] = 0
rcnn_boxes3d = self.box_coder.decode_torch(
fg_rcnn_reg.view(batch_anchors.shape[0], -1, code_size), batch_anchors
).view(-1, code_size)
rcnn_boxes3d = common_utils.rotate_points_along_z(
rcnn_boxes3d.unsqueeze(dim=1), roi_ry
).squeeze(dim=1)
rcnn_boxes3d[:, 0:3] += roi_xyz
loss_corner = loss_utils.get_corner_loss_lidar_v1(
rcnn_boxes3d[:, 0:7],
gt_of_rois_src[fg_mask][:, 0:7]
)
loss_corner = loss_corner.mean()
loss_corner = loss_corner * loss_cfgs.LOSS_WEIGHTS['rcnn_corner_weight']
rcnn_loss_reg += loss_corner
tb_dict['rcnn_loss_corner'] = loss_corner.item()
else:
raise NotImplementedError
return rcnn_loss_reg, tb_dict
