2. 数据增强策略实现
任务说明
本案例实现端到端自动驾驶模型训练中常用的数据增强策略,包括随机水平翻转、随机旋转和颜色抖动。目标是通过在训练时对输入图像施加受控扰动,提升模型对光照变化、车辆姿态偏移和左右对称场景的泛化能力。代码将提供一个可独立运行的增强模块,并给出单张图像的调用示例与批量增强的验证方式。
实现方案
采用OpenCV与NumPy实现全部增强操作,不依赖深度学习框架,便于独立验证。随机水平翻转模拟车辆左右对称的驾驶场景;随机旋转模拟车辆在坡道或颠簸路面上出现的俯仰/侧倾视角变化;颜色抖动通过调整HSV空间的饱和度与明度,模拟不同天气和光照条件。所有增强均以随机概率触发,且旋转与颜色抖动幅度限制在合理范围内,避免破坏驾驶语义信息。
代码实现
import cv2
import numpy as np
class DrivingDataAugmentor:
"""端到端自动驾驶训练数据增强器"""
def __init__(self, flip_prob=0.5, max_rotation=10.0,
saturation_range=(0.75, 1.25), value_range=(0.75, 1.25)):
self.flip_prob = flip_prob
self.max_rotation = max_rotation
self.saturation_range = saturation_range
self.value_range = value_range
def random_flip(self, image):
"""随机水平翻转(模拟左右对称驾驶场景)"""
if np.random.random() < self.flip_prob:
return cv2.flip(image, 1)
return image
def random_rotation(self, image):
"""随机旋转(模拟坡道/颠簸路面视角变化)"""
if np.random.random() < self.flip_prob:
h, w = image.shape[:2]
angle = np.random.uniform(-self.max_rotation, self.max_rotation)
matrix = cv2.getRotationMatrix2D((w/2, h/2), angle, 1.0)
return cv2.warpAffine(image, matrix, (w, h),
flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_REFLECT)
return image
def color_jitter(self, image):
"""颜色抖动(调整饱和度与明度,模拟光照变化)"""
if np.random.random() < self.flip_prob:
hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV).astype(np.float32)
sat_scale = np.random.uniform(*self.saturation_range)
val_scale = np.random.uniform(*self.value_range)
hsv[:, :, 1] = np.clip(hsv[:, :, 1] * sat_scale, 0, 255)
hsv[:, :, 2] = np.clip(hsv[:, :, 2] * val_scale, 0, 255)
return cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
return image
def __call__(self, image):
"""依次执行翻转→旋转→颜色抖动"""
image = self.random_flip(image)
image = self.random_rotation(image)
image = self.color_jitter(image)
return image
# 独立验证示例
if __name__ == "__main__":
# 生成一张模拟驾驶图像(640x480 灰度渐变+矩形障碍物)
demo_img = np.zeros((480, 640, 3), dtype=np.uint8)
demo_img[:, :, 0] = np.linspace(80, 200, 640, dtype=np.uint8) # B通道渐变
demo_img[:, :, 1] = 100 # G通道固定
demo_img[:, :, 2] = np.linspace(200, 80, 640, dtype=np.uint8) # R通道渐变
cv2.rectangle(demo_img, (200, 150), (350, 300), (255, 255, 255), -1) # 白色障碍物
augmentor = DrivingDataAugmentor()
# 固定随机种子,保证可复现
np.random.seed(42)
augmented = augmentor(demo_img)
# 验证输出尺寸与数据类型
assert augmented.shape == demo_img.shape, "输出尺寸必须与输入一致"
assert augmented.dtype == np.uint8, "输出类型必须为uint8"
# 验证翻转确实发生(白色矩形中心x坐标应镜像)
orig_center = np.where(demo_img[:, :, 0] > 250)
aug_center = np.where(augmented[:, :, 0] > 250)
if len(orig_center[1]) > 0 and len(aug_center[1]) > 0:
orig_x = int(np.mean(orig_center[1]))
aug_x = int(np.mean(aug_center[1]))
print(f"翻转验证: 原始中心x={orig_x}, 增强后中心x={aug_x}, 镜像关系={abs(orig_x + aug_x - 639) < 5}")
# 保存结果供人工检查
cv2.imwrite("original.jpg", demo_img)
cv2.imwrite("augmented.jpg", augmented)
print("验证完成,已保存 original.jpg 与 augmented.jpg")
落地与验证
运行方式:将上述代码保存为 data_augment.py,执行 python data_augment.py。代码会生成一张模拟驾驶图像并施加增强,输出两张图片供对比。
验证方法:程序内置三项自动检查——输出尺寸与输入一致、数据类型为uint8、翻转后白色障碍物的水平中心坐标满足镜像关系(即原中心x与增强后中心x之和接近图像宽度减1)。若三项均通过,终端会打印翻转验证信息并显示保存成功。
判定标准:运行无异常且打印信息显示镜像关系成立,同时人工查看 augmented.jpg 可见图像发生水平翻转、轻微旋转和颜色变化,即判定本案例实现正确。若需用于模型训练,在训练循环中调用 augmentor(batch_image) 即可,每次调用会按概率独立决定是否执行各增强操作。