3. BEV与Occupancy数据增强策略实现
任务目标
实现一套同时作用于BEV栅格特征与3D Occupancy标签的数据增强方法，确保几何变换（翻转、旋转、缩放）与颜色抖动在两类数据上严格一致，避免联合训练时出现空间错位。
核心实现
import numpy as np
import random
from scipy.ndimage import rotate, zoom
class BEVOccAugmentor:
def __init__(self, bev_size=(200, 200), occ_size=(200, 200, 16),
flip_prob=0.5, rot_range=(-5, 5), scale_range=(0.95, 1.05),
color_jitter=(0.2, 0.2, 0.2)):
self.bev_size = bev_size
self.occ_size = occ_size
self.flip_prob = flip_prob
self.rot_range = rot_range
self.scale_range = scale_range
self.color_jitter = color_jitter
def _get_transform_params(self):
"""生成一组统一的几何变换参数"""
do_flip = random.random() < self.flip_prob
angle = random.uniform(*self.rot_range)
scale = random.uniform(*self.scale_range)
return do_flip, angle, scale
def _apply_geometry(self, data, do_flip, angle, scale, is_occ=False):
"""对BEV或Occupancy应用相同的几何变换"""
if is_occ:
# Occupancy: (H, W, D) 逐通道旋转/缩放
transformed = np.zeros_like(data)
for d in range(data.shape[2]):
slice_2d = data[:, :, d]
if do_flip:
slice_2d = np.fliplr(slice_2d)
if angle != 0:
slice_2d = rotate(slice_2d, angle, reshape=False, order=0)
if scale != 1.0:
zoomed = zoom(slice_2d, scale, order=0)
# 裁剪或填充回原尺寸
h, w = slice_2d.shape
zh, zw = zoomed.shape
h_start = max(0, (zh - h) // 2)
w_start = max(0, (zw - w) // 2)
slice_2d = zoomed[h_start:h_start+h, w_start:w_start+w]
if slice_2d.shape != (h, w):
slice_2d = np.pad(slice_2d,
((0, h - slice_2d.shape[0]),
(0, w - slice_2d.shape[1])),
mode='constant')
transformed[:, :, d] = slice_2d
return transformed
else:
# BEV: (H, W, C) 所有通道同步变换
transformed = np.zeros_like(data)
for c in range(data.shape[2]):
channel = data[:, :, c]
if do_flip:
channel = np.fliplr(channel)
if angle != 0:
channel = rotate(channel, angle, reshape=False, order=0)
if scale != 1.0:
zoomed = zoom(channel, scale, order=0)
h, w = channel.shape
zh, zw = zoomed.shape
h_start = max(0, (zh - h) // 2)
w_start = max(0, (zw - w) // 2)
channel = zoomed[h_start:h_start+h, w_start:w_start+w]
if channel.shape != (h, w):
channel = np.pad(channel,
((0, h - channel.shape[0]),
(0, w - channel.shape[1])),
mode='constant')
transformed[:, :, c] = channel
return transformed
def _apply_color_jitter(self, bev):
"""仅对BEV图像做颜色抖动（Occupancy标签不做）"""
jittered = bev.copy()
for c in range(bev.shape[2]):
brightness = 1.0 + random.uniform(-self.color_jitter[0],
self.color_jitter[0])
contrast = 1.0 + random.uniform(-self.color_jitter[1],
self.color_jitter[1])
saturation = 1.0 + random.uniform(-self.color_jitter[2],
self.color_jitter[2])
jittered[:, :, c] = np.clip(
(bev[:, :, c] * brightness - 0.5) * contrast + 0.5,
0, 1) * saturation
return np.clip(jittered, 0, 1)
def __call__(self, bev, occupancy):
"""
bev: (H, W, C) float32, 值域[0,1]
occupancy: (H, W, D) int, 0/1 体素占用
返回增强后的 (bev_aug, occ_aug)
"""
assert bev.shape[:2] == self.bev_size, \
f"BEV尺寸应为{self.bev_size}, 实际{bev.shape[:2]}"
assert occupancy.shape[:2] == self.occ_size[:2], \
f"Occupancy前两维应为{self.occ_size[:2]}, 实际{occupancy.shape[:2]}"
do_flip, angle, scale = self._get_transform_params()
# 几何变换：BEV与Occupancy使用同一组参数
bev_geo = self._apply_geometry(bev, do_flip, angle, scale,
is_occ=False)
occ_geo = self._apply_geometry(occupancy, do_flip, angle, scale,
is_occ=True)
# 颜色抖动：仅BEV
bev_aug = self._apply_color_jitter(bev_geo)
return bev_aug, occ_geo
# 最小验证闭环
if __name__ == "__main__":
# 构造确定性输入
np.random.seed(42)
bev = np.random.rand(200, 200, 3).astype(np.float32)
occ = np.random.randint(0, 2, (200, 200, 16)).astype(np.int8)
aug = BEVOccAugmentor(flip_prob=1.0, rot_range=(0, 0),
scale_range=(1.0, 1.0),
color_jitter=(0, 0, 0))
bev_aug, occ_aug = aug(bev, occ)
# 验证1：翻转后BEV与Occupancy空间对齐
# 取一个已知位置的特征点，翻转后应出现在镜像位置
test_bev = np.zeros((200, 200, 1), dtype=np.float32)
test_bev[50, 80, 0] = 1.0
test_occ = np.zeros((200, 200, 16), dtype=np.int8)
test_occ[50, 80, 5] = 1
aug2 = BEVOccAugmentor(flip_prob=1.0, rot_range=(0, 0),
scale_range=(1.0, 1.0),
color_jitter=(0, 0, 0))
bev_t, occ_t = aug2(test_bev, test_occ)
# 翻转后x坐标变为 200-1-80 = 119
assert bev_t[50, 119, 0] == 1.0, "BEV翻转位置错误"
assert occ_t[50, 119, 5] == 1, "Occupancy翻转位置错误"
print("翻转一致性验证通过")
# 验证2：旋转一致性（旋转90度）
aug3 = BEVOccAugmentor(flip_prob=0.0, rot_range=(90, 90),
scale_range=(1.0, 1.0),
color_jitter=(0, 0, 0))
test_bev2 = np.zeros((200, 200, 1), dtype=np.float32)
test_bev2[50, 80, 0] = 1.0
test_occ2 = np.zeros((200, 200, 16), dtype=np.int8)
test_occ2[50, 80, 5] = 1
bev_r, occ_r = aug3(test_bev2, test_occ2)
# 旋转90度后 (50,80) -> (80, 149)
assert bev_r[80, 149, 0] == 1.0, "BEV旋转位置错误"
assert occ_r[80, 149, 5] == 1, "Occupancy旋转位置错误"
print("旋转一致性验证通过")
# 验证3：颜色抖动不影响Occupancy
aug4 = BEVOccAugmentor(flip_prob=0.0, rot_range=(0, 0),
scale_range=(1.0, 1.0),
color_jitter=(0.5, 0.5, 0.5))
bev_c, occ_c = aug4(bev, occ)
assert np.array_equal(occ_c, occ), "颜色抖动不应改变Occupancy"
assert not np.array_equal(bev_c, bev), "颜色抖动应改变BEV"
print("颜色抖动隔离验证通过")运行与验证
依赖：numpy、scipy（pip install numpy scipy）
运行：直接执行脚本 python bev_occ_augment.py，若输出以下三行则通过：
翻转一致性验证通过
旋转一致性验证通过
颜色抖动隔离验证通过关键设计说明
- 参数共享：几何变换参数（翻转、旋转角、缩放比）由
_get_transform_params一次性生成，BEV与Occupancy共用，从根源上保证空间一致性。 - Occupancy逐层处理：3D体素标签按高度维逐层应用2D几何变换，避免3D插值引入伪占用。
- 缩放边界处理：缩放后裁剪或填充回原始尺寸，保证输出形状不变，便于直接送入网络。
- 颜色抖动隔离：仅作用于BEV图像通道，Occupancy标签不受影响，符合语义（颜色扰动不改变占用状态）。