2. 数据增强策略实现

任务说明

本案例实现端到端自动驾驶模型训练中常用的数据增强策略,包括随机水平翻转、随机旋转和颜色抖动。目标是通过在训练时对输入图像施加受控扰动,提升模型对光照变化、车辆姿态偏移和左右对称场景的泛化能力。代码将提供一个可独立运行的增强模块,并给出单张图像的调用示例与批量增强的验证方式。

实现方案

采用OpenCV与NumPy实现全部增强操作,不依赖深度学习框架,便于独立验证。随机水平翻转模拟车辆左右对称的驾驶场景;随机旋转模拟车辆在坡道或颠簸路面上出现的俯仰/侧倾视角变化;颜色抖动通过调整HSV空间的饱和度与明度,模拟不同天气和光照条件。所有增强均以随机概率触发,且旋转与颜色抖动幅度限制在合理范围内,避免破坏驾驶语义信息。

原始图像 (H,W,3) uint8 随机水平翻转 p=0.5 随机旋转 ±10° 颜色抖动 饱和度/明度 ±25% p=0.5 执行 p=0.5 执行 p=0.5 执行 增强后图像 (H,W,3) uint8 每次调用按概率独立决定是否执行各增强操作

代码实现


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) 即可,每次调用会按概率独立决定是否执行各增强操作。