点云数据增强方法

目录

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

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

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

四、过滤掉范围外的点云或获取GT内的点云

五、随机缩放 

六、全局沿xyz轴平移 

 七、GT沿xyz平移

八、对GT缩放

九、对GT旋转 


一、全局沿着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 global_rotation(gt_boxes, points, rot_range):
    """
    对点云和box进行整体旋转
    Args:
        gt_boxes: (N, 7 + C), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
        points: (M, 3 + C),
        rot_range: [min, max]
    Returns:
    """
    # 在均匀分布中随机产生旋转角度
    noise_rotation = np.random.uniform(rot_range[0], rot_range[1])
    # 沿z轴旋转noise_rotation弧度,这里之所以取第0个,是因为rotate_points_along_z对batch进行处理,而这里仅处理单个点云
    #np.newaxis 放在哪个位置,就会给哪个位置增加维度
    points = common_utils.rotate_points_along_z(points[np.newaxis, :, :], np.array([noise_rotation]))[0]
    # 同样对box的坐标进行旋转
    gt_boxes[:, 0:3] = common_utils.rotate_points_along_z(gt_boxes[np.newaxis, :, 0:3], np.array([noise_rotation]))[0]
    # 对box的方位角进行累加
    gt_boxes[:, 6] += noise_rotation

    return gt_boxes, points

四、过滤掉范围外的点云或获取GT内的点云

def mask_points_by_range(points, limit_range):
    # 根据点云的范围产生mask,过滤点云
    mask = (points[:, 0] >= limit_range[0]) & (points[:, 0] <= limit_range[3]) \
           & (points[:, 1] >= limit_range[1]) & (points[:, 1] <= limit_range[4])
    return mask
def get_points_in_box(points, gt_box):
    x, y, z = points[:, 0], points[:, 1], points[:, 2]
    cx, cy, cz = gt_box[0], gt_box[1], gt_box[2]
    dx, dy, dz, rz = gt_box[3], gt_box[4], gt_box[5], gt_box[6]
    shift_x, shift_y, shift_z = x - cx, y - cy, z - cz

    MARGIN = 1e-1
    cosa, sina = math.cos(-rz), math.sin(-rz)
    local_x = shift_x * cosa + shift_y * (-sina)
    local_y = shift_x * sina + shift_y * cosa

    mask = np.logical_and(abs(shift_z) <= dz / 2.0, \
                          np.logical_and(abs(local_x) <= dx / 2.0 + MARGIN, \
                                         abs(local_y) <= dy / 2.0 + MARGIN))

    points = points[mask]

    return points, mask

五、随机缩放 

points和GT直接乘缩放系数

def global_scaling(gt_boxes, points, scale_range):
    """
    随机缩放
    Args:
        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading]
        points: (M, 3 + C),
        scale_range: [min, max]
    Returns:
    """
    # 如果缩放的尺度过小,则直接返回原来的box和点云
    if scale_range[1] - scale_range[0] < 1e-3:
        return gt_boxes, points
    # 在缩放范围内随机产生缩放尺度
    noise_scale = np.random.uniform(scale_range[0], scale_range[1])
    # 将点云和box同时乘以缩放尺度
    points[:, :3] *= noise_scale
    gt_boxes[:, :6] *= noise_scale

    return gt_boxes, points

六、全局沿xyz轴平移 

# 沿X轴随机平移
def random_translation_along_x(gt_boxes, points, offset_range):
    """
    Args:
        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
        points: (M, 3 + C),
        offset_range: [min max]]
    Returns:
    """
    offset = np.random.uniform(offset_range[0], offset_range[1])

    points[:, 0] += offset
    gt_boxes[:, 0] += offset

    # if gt_boxes.shape[1] > 7:
    #     gt_boxes[:, 7] += offset

    return gt_boxes, points


def random_translation_along_y(gt_boxes, points, offset_range):
    """
    Args:
        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
        points: (M, 3 + C),
        offset_range: [min max]]
    Returns:
    """
    offset = np.random.uniform(offset_range[0], offset_range[1])

    points[:, 1] += offset
    gt_boxes[:, 1] += offset

    # if gt_boxes.shape[1] > 8:
    #     gt_boxes[:, 8] += offset

    return gt_boxes, points


def random_translation_along_z(gt_boxes, points, offset_range):
    """
    Args:
        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
        points: (M, 3 + C),
        offset_range: [min max]]
    Returns:
    """
    offset = np.random.uniform(offset_range[0], offset_range[1])

    points[:, 2] += offset
    gt_boxes[:, 2] += offset

    return gt_boxes, points

 七、GT沿xyz平移

先找到GT内的点,对这些点和GT中心平移

def random_local_translation_along_x(gt_boxes, points, offset_range):
    """
    Args:
        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
        points: (M, 3 + C),
        offset_range: [min max]]
    Returns:
    """
    # augs = {}
    for idx, box in enumerate(gt_boxes):
        offset = np.random.uniform(offset_range[0], offset_range[1])
        # augs[f'object_{idx}'] = offset
        points_in_box, mask = get_points_in_box(points, box)
        points[mask, 0] += offset

        gt_boxes[idx, 0] += offset

        # if gt_boxes.shape[1] > 7:
        #     gt_boxes[idx, 7] += offset

    return gt_boxes, points


def random_local_translation_along_y(gt_boxes, points, offset_range):
    """
    Args:
        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
        points: (M, 3 + C),
        offset_range: [min max]]
    Returns:
    """
    # augs = {}
    for idx, box in enumerate(gt_boxes):
        offset = np.random.uniform(offset_range[0], offset_range[1])
        # augs[f'object_{idx}'] = offset
        points_in_box, mask = get_points_in_box(points, box)
        points[mask, 1] += offset

        gt_boxes[idx, 1] += offset

        # if gt_boxes.shape[1] > 8:
        #     gt_boxes[idx, 8] += offset

    return gt_boxes, points


def random_local_translation_along_z(gt_boxes, points, offset_range):
    """
    Args:
        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
        points: (M, 3 + C),
        offset_range: [min max]]
    Returns:
    """
    # augs = {}
    for idx, box in enumerate(gt_boxes):
        offset = np.random.uniform(offset_range[0], offset_range[1])
        # augs[f'object_{idx}'] = offset
        points_in_box, mask = get_points_in_box(points, box)
        points[mask, 2] += offset

        gt_boxes[idx, 2] += offset

    return gt_boxes, points

八、对GT缩放

得到GT内的点,然后和GT中心做差,乘系数,再加回来

def local_scaling(gt_boxes, points, scale_range):
    """
    Args:
        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading]
        points: (M, 3 + C),
        scale_range: [min, max]
    Returns:
    """
    if scale_range[1] - scale_range[0] < 1e-3:
        return gt_boxes, points

    # augs = {}
    for idx, box in enumerate(gt_boxes):
        noise_scale = np.random.uniform(scale_range[0], scale_range[1])
        # augs[f'object_{idx}'] = noise_scale
        points_in_box, mask = get_points_in_box(points, box)

        # tranlation to axis center
        points[mask, 0] -= box[0]
        points[mask, 1] -= box[1]
        points[mask, 2] -= box[2]

        # apply scaling
        points[mask, :3] *= noise_scale

        # tranlation back to original position
        points[mask, 0] += box[0]
        points[mask, 1] += box[1]
        points[mask, 2] += box[2]

        gt_boxes[idx, 3:6] *= noise_scale
    return gt_boxes, points

九、对GT旋转 

得到GT里面的点,减去中心归一化,乘旋转矩阵,然后再加回来中心

def local_rotation(gt_boxes, points, rot_range):
    """
    Args:
        gt_boxes: (N, 7), [x, y, z, dx, dy, dz, heading, [vx], [vy]]
        points: (M, 3 + C),
        rot_range: [min, max]
    Returns:
    """
    # augs = {}
    for idx, box in enumerate(gt_boxes):
        noise_rotation = np.random.uniform(rot_range[0], rot_range[1])
        # augs[f'object_{idx}'] = noise_rotation
        points_in_box, mask = get_points_in_box(points, box)

        centroid_x = box[0]
        centroid_y = box[1]
        centroid_z = box[2]

        # tranlation to axis center
        points[mask, 0] -= centroid_x
        points[mask, 1] -= centroid_y
        points[mask, 2] -= centroid_z
        box[0] -= centroid_x
        box[1] -= centroid_y
        box[2] -= centroid_z

        # apply rotation
        points[mask, :] = common_utils.rotate_points_along_z(points[np.newaxis, mask, :], np.array([noise_rotation]))[0]
        box[0:3] = common_utils.rotate_points_along_z(box[np.newaxis, np.newaxis, 0:3], np.array([noise_rotation]))[0][
            0]

        # tranlation back to original position
        points[mask, 0] += centroid_x
        points[mask, 1] += centroid_y
        points[mask, 2] += centroid_z
        box[0] += centroid_x
        box[1] += centroid_y
        box[2] += centroid_z

        gt_boxes[idx, 6] += noise_rotation

    return gt_boxes, points

你可能感兴趣的:(python,开发语言)