您现在的位置是:首页 >技术杂谈 >MMDet3D——数据增强Pipline‘GlobalRotScaleTrans‘和‘RandomFlip3D‘的Pytorch逆变换实现网站首页技术杂谈

MMDet3D——数据增强Pipline‘GlobalRotScaleTrans‘和‘RandomFlip3D‘的Pytorch逆变换实现

Irving.Gao 2024-08-19 00:01:04
简介MMDet3D——数据增强Pipline‘GlobalRotScaleTrans‘和‘RandomFlip3D‘的Pytorch逆变换实现

在点云的3D感知算法中,常用RandomFlip3DGlobalRotScaleTrans的数据增强方式,这两个可以有效地增强模型的鲁棒性,提升模型的性能。

		transforms=[
            dict(
                type='RandomFlip3D',
                sync_2d=False,
                flip_ratio_bev_horizontal=0.5,
                flip_ratio_bev_vertical=0.5),
            dict(
                type='GlobalRotScaleTrans',
                rot_range=[-0.78539816, 0.78539816],
                scale_ratio_range=[0.95, 1.05]),

而本文的出发点在于当我们同时对相同的data使用了多种不同的随机增强方式时,当我们需要在模型中将不同增强方法下的Feature map或者GT 3D Box进行对齐时,就需要涉及到逆变换。
尤其适用于半监督等同时有label和unlabel data的代码。

简单粗暴,我们直接上代码,注释直接写在代码中:(代码可以直接copy使用)

	def align_aug_dataV2(self, pts_feats, # 可以是feature map, shape为 [B,C,H,W]
                         gt_bboxes_3d,# Nuscenes默认的LidarBox格式, shape为 [B,N,x], x=7或9,有无速度的区别
                         img_metas,
                         return_tensor=True,			# 如果想返回的pts_feats是tensor形式,默认True
                         interploate_mode='bilinear'):	# grid_sample的插值方式,如果是feat map,建议bilinear
        '''
        目前的box.tensor的格式为[x, y, z, x_size, y_size, z_size, yaw, vx, vy],shape为[N, 9]
        所以对gt_box做flip需要分别处理x,y,yaw和vx,vy
        '''
        # feature_map: [B,C,H,W]
        def horizontal_flip(feature_map):# 水平翻转
            return torch.flip(feature_map, [2])
        def vertical_flip(feature_map):# 垂直翻转
            return torch.flip(feature_map, [1])
        def box_flip(boxes, bev_direction='horizontal'):
            assert bev_direction in ('horizontal', 'vertical')
            if bev_direction == 'horizontal':
                boxes.tensor[:, 1] = -boxes.tensor[:, 1]    # y
                boxes.tensor[:, 6] = -boxes.tensor[:, 6] + np.pi 
                boxes.tensor[:, 7] = boxes.tensor[:, 7]    # vx
                boxes.tensor[:, 8] = -boxes.tensor[:, 8]    # vy
            elif bev_direction == 'vertical':
                boxes.tensor[:, 0] = -boxes.tensor[:, 0]    # x
                boxes.tensor[:, 6] = -boxes.tensor[:, 6]
                boxes.tensor[:, 7] = -boxes.tensor[:, 7]    # vx
                boxes.tensor[:, 8] = boxes.tensor[:, 8]    # vy
            return boxes
        aligned_pts_feats = []
        aligned_gt_bboxes_3d = deepcopy(gt_bboxes_3d)
        for idx, (pts_feat, boxes_3d, img_meta) in enumerate(
            zip(pts_feats, aligned_gt_bboxes_3d, img_metas)):
            # 数据增强Pipline处理时aug顺序: RandomFlip3D, GlobalRotScaleTrans(rot, scale, trans)
            # 逆变换align顺序需要倒过来
            # ----------------------------------------------------
            pts_feat = pts_feat.unsqueeze(0)
            tgt_size = pts_feat.shape
            dev = pts_feat.device
            # GlobalRotScaleTrans
            if 'pcd_trans' in img_meta:
                if not (img_meta['pcd_trans'] == 0.).all():
                	# TODO: 这一部分我没有测试,因为默认没有使用到,所有有需要自行进行测试
                    # 1. feat map
                    Trans = torch.zeros_like(img_meta['pcd_rotation'].T)    # [3,3]
                    Trans[0,0], Trans[1,1] = 1,1
                    Trans[:2,2] = img_meta['pcd_trans'][:2]
                    Trans = Trans[:2,:].unsqueeze(0)        # [B,2,3]
                    grid = F.affine_grid(Trans, tgt_size).to(dev)   # # 仿射变换矩阵
                    pts_feat = F.grid_sample(pts_feat, # 输入tensor,shape为[B,C,W,H]
                                        grid, # 上一步输出的gird,shape为[B,C,W,H]
                                        mode=interploate_mode)
                    # 2. gt_boxes
                    boxes_3d.translate(-img_meta['pcd_trans'])
            
            if 'pcd_scale_factor' in img_meta:
                if img_meta['pcd_scale_factor'] != 1.:
                    # 1. feat map
                    Scl = torch.zeros_like(img_meta['pcd_rotation'].T)    # [3,3]
                    Scl[0,0], Scl[1,1] = 1/img_meta['pcd_scale_factor'], 1/img_meta['pcd_scale_factor']
                    Scl = Scl[:2,:].unsqueeze(0)        # [B,2,3]
                    grid = F.affine_grid(Scl, tgt_size).to(dev)   # # 仿射变换矩阵
                    pts_feat = F.grid_sample(pts_feat, # 输入tensor,shape为[B,C,W,H]
                                        grid, # 上一步输出的gird,shape为[B,C,W,H]
                                        mode=interploate_mode)
                    # 2. gt_boxes
                    boxes_3d.scale(1/img_meta['pcd_scale_factor'])
                    
            if 'pcd_rotation' in img_meta:
                if img_meta['pcd_rotation'][0,0] != 1.:
                    # 1. feat map
                    Rot = img_meta['pcd_rotation'].T    # [3,3]
                    Rot = Rot[:2,:].unsqueeze(0)        # [B,2,3]
                    grid = F.affine_grid(Rot, tgt_size).to(dev)   # # 仿射变换矩阵
                    pts_feat = F.grid_sample(pts_feat, # 输入tensor,shape为[B,C,W,H]
                                        grid, # 上一步输出的gird,shape为[B,C,W,H]
                                        mode=interploate_mode)
                    # 2. gt_boxes
                    boxes_3d.rotate(img_meta['pcd_rotation'].T)
            pts_feat = pts_feat[0]
            
            # ----------------------------------------------------
            # RandomFlip3D
            if 'pcd_vertical_flip' in img_meta:
                if img_meta['pcd_vertical_flip']:
                    # 1. feat map
                    pts_feat = vertical_flip(pts_feat)
                    # 2. gt_boxes
                    boxes_3d = box_flip(boxes_3d, bev_direction='vertical')

            if 'pcd_horizontal_flip' in img_meta:
                if img_meta['pcd_horizontal_flip']:
                    # 1. feat map
                    pts_feat = horizontal_flip(pts_feat)
                    # 2. gt_boxes
                    # boxes_3d.flip(bev_direction='horizontal')
                    boxes_3d = box_flip(boxes_3d, bev_direction='horizontal')
                    
            aligned_pts_feats.append(pts_feat)
            aligned_gt_bboxes_3d[idx] = boxes_3d
        if return_tensor:
            aligned_pts_feats = torch.stack(aligned_pts_feats)

        return aligned_pts_feats, aligned_gt_bboxes_3d

有几点需要注意:

  • feat map在经过Flip之后也不全是一一对应,因为在卷积过程中,翻转后对应的位置不同,所以feat在经过翻转变换之后略有不同是正常现象;

测试效果

测试一下把有aug的feat map和GT都 对齐到没有aug的情况下:

# 测试代码: (1-3应当相同)
# 1.aligned_aug     aligned_gt_bboxes_3d[0].tensor[0]
# 2.aligned_aug     aligned_gt_bboxes_3d[4].tensor[0]
# 3.无aug           gt_bboxes_3d[4].tensor[0]
# 4.有aug           gt_bboxes_3d[0].tensor[0]
  • 加上RandomFlip3D:(标红的为有aug的原始值
    在这里插入图片描述

  • 加上RandomFlip3DScale
    在这里插入图片描述

  • 加上RandomFlip3DRotate
    在这里插入图片描述

  • 加上RandomFlip3DGlobalRotScaleTrans所有aug的对齐结果:
    在这里插入图片描述

参考文章:

风语者!平时喜欢研究各种技术,目前在从事后端开发工作,热爱生活、热爱工作。