点云数据增强方法
目录
一、全局沿着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 maskdef 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