点云数据增强方法

原创

已于 2022-11-02 10:19:39 修改 · 2.5k 阅读

34 ·

CC 4.0 BY-SA版权

文章标签：

#python #开发语言

于 2022-11-01 20:24:14 首次发布

本文介绍了针对点云数据的多种增强方法，包括沿不同轴的翻转和平移、全局及局部的旋转与缩放等操作，旨在提升点云数据集的质量和模型训练效果。

一、全局沿着x轴随机翻转

x坐标不变，y取反，gt的方位角取反

def random_flip_along_x(gt_boxes, points):
    """
    沿着x轴随机翻转
    Args:
        gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
        points: (M, 3 + C)
    Returns:
    """
    # 随机选择是否翻转      replace=False表示不可以取相同数字，p对应前面的概率
    enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5])
    if enable:
        #沿着x轴随机翻转，x坐标不变，y取反，同时方位角定义的是与x轴的夹角
        gt_boxes[:, 1] = -gt_boxes[:, 1]  # y坐标翻转
        gt_boxes[:, 6] = -gt_boxes[:, 6]  # 方位角翻转，直接取负数，因为方位角定义为与x轴的夹角（这里按照顺时针的方向取角度）
        points[:, 1] = -points[:, 1]  # 点云y坐标翻转


    return gt_boxes, points

二、全局沿着y轴随机翻转

y坐标不变，x取反，GT方位角加Π取反

def random_flip_along_y(gt_boxes, points):
    """
    沿着y轴随机翻转
    Args:
        gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
        points: (M, 3 + C)
    Returns:
    """
    # 随机旋转是否翻转
    enable = np.random.choice([False, True], replace=False, p=[0.5, 0.5])
    if enable:
        gt_boxes[:, 0] = -gt_boxes[:, 0]  # x坐标翻转
        gt_boxes[:, 6] = -(gt_boxes[:, 6] + np.pi)  # 方位角加pi后，取负数（这里按照顺时针的方向取角度）
        points[:, 0] = -points[:, 0]  # 点云x坐标取反

    return gt_boxes, points

三、全局沿着z轴随机旋转

先判断是不是tensor张量，不是就把numpy转tensor，构建旋转矩阵，torch.stack，再相乘，GT角度直接加

def check_numpy_to_torch(x):
    # 检测输入数据是否是numpy格式，如果是，则转换为torch格式
    if isinstance(x, np.ndarray):
        return torch.from_numpy(x).float(), True
    return x, False

def rotate_points_along_z(points, angle):
    """
    Args:
        points: (B, N, 3 + C)
        angle: (B), angle along z-axis, angle increases x ==> y
    Returns:

    """
    # 首先利用torch.from_numpy().float将numpy转化为torch
    points, is_numpy = check_numpy_to_torch(points)
    angle, _ = check_numpy_to_torch(angle)

    # 构造旋转矩阵batch个
    cosa = torch.cos(angle)
    sina = torch.sin(angle)
    zeros = angle.new_zeros(points.shape[0])
    ones = angle.new_ones(points.shape[0])
    rot_matrix = torch.stack((
        cosa,  sina, zeros,
        -sina, cosa, zeros,
        zeros, zeros, ones
    ), dim=1).view(-1, 3, 3).float()
    # 对点云坐标进行旋转
    points_rot = torch.matmul(points[:, :, 0:3], rot_matrix)
    # 将旋转后的点云坐标与原始额外特征拼接
    points_rot = torch.cat((points_rot, points[:, :, 3:]), dim=-1)
    # 将点云转化为numpy格式，并返回
    return points_rot.numpy() if is_numpy else points_rot

def glob