2. 图像增强与数据扩充

任务目标

端到端自动驾驶模仿学习中,模型常因训练数据分布单一而过拟合。本案例对已采集的驾驶图像(前视摄像头画面)应用随机亮度、对比度、平移、旋转等增强操作,生成多样化的扩充样本,提升模型对光照变化与视角偏移的泛化能力。

环境依赖

需要 opencv-python、numpy。若需可视化增强结果,额外安装 matplotlib(可选)。

关键完整代码

import cv2
import numpy as np
import os
import random

def apply_brightness_contrast(img, brightness_delta=40, contrast_alpha=(0.8, 1.2)):
    """随机调整亮度与对比度"""
    delta = random.uniform(-brightness_delta, brightness_delta)
    alpha = random.uniform(*contrast_alpha)
    adjusted = cv2.convertScaleAbs(img, alpha=alpha, beta=delta)
    return adjusted

def apply_translation(img, max_shift_ratio=0.08):
    """随机平移(水平/垂直方向)"""
    h, w = img.shape[:2]
    max_dx = int(w * max_shift_ratio)
    max_dy = int(h * max_shift_ratio)
    dx = random.randint(-max_dx, max_dx)
    dy = random.randint(-max_dy, max_dy)
    M = np.float32([[1, 0, dx], [0, 1, dy]])
    shifted = cv2.warpAffine(img, M, (w, h), borderMode=cv2.BORDER_REPLICATE)
    return shifted

def apply_rotation(img, max_angle=10):
    """随机旋转(角度范围 ±max_angle 度)"""
    h, w = img.shape[:2]
    angle = random.uniform(-max_angle, max_angle)
    M = cv2.getRotationMatrix2D((w // 2, h // 2), angle, 1.0)
    rotated = cv2.warpAffine(img, M, (w, h), borderMode=cv2.BORDER_REPLICATE)
    return rotated

def augment_image(img):
    """对单张图像依次应用随机增强(每次调用产生不同组合)"""
    # 随机顺序执行增强,避免固定模式
    ops = [
        lambda x: apply_brightness_contrast(x),
        lambda x: apply_translation(x),
        lambda x: apply_rotation(x)
    ]
    random.shuffle(ops)
    result = img.copy()
    for op in ops:
        result = op(result)
    return result

def augment_dataset(input_dir, output_dir, copies_per_image=5, seed=42):
    """
    对 input_dir 下所有 .jpg/.png 图像进行扩充。
    每张原图生成 copies_per_image 张增强图,保存到 output_dir。
    """
    random.seed(seed)
    np.random.seed(seed)
    os.makedirs(output_dir, exist_ok=True)

    valid_exts = ('.jpg', '.jpeg', '.png')
    image_files = [f for f in os.listdir(input_dir) if f.lower().endswith(valid_exts)]
    if not image_files:
        raise ValueError(f"输入目录 {input_dir} 中未找到图像文件")

    for fname in image_files:
        img_path = os.path.join(input_dir, fname)
        img = cv2.imread(img_path)
        if img is None:
            print(f"警告:无法读取 {img_path},跳过")
            continue

        base_name = os.path.splitext(fname)[0]
        for i in range(copies_per_image):
            aug_img = augment_image(img)
            out_path = os.path.join(output_dir, f"{base_name}_aug{i}.jpg")
            cv2.imwrite(out_path, aug_img)

    print(f"完成:{len(image_files)} 张原图,每张生成 {copies_per_image} 张增强图,共 {len(image_files) * copies_per_image} 张,已保存至 {output_dir}")

if __name__ == "__main__":
    # 使用示例:将 "raw_images" 目录中的图像扩充到 "augmented_images"
    augment_dataset("raw_images", "augmented_images", copies_per_image=5)

运行方式

  1. 准备一个存放原始图像的目录(如 raw_images/),放入若干 .jpg 或 .png 前视驾驶图像。
  2. 将上述代码保存为 augment.py,在终端执行:
python augment.py

验证方法

  • 文件数量检查:运行后 augmented_images/ 中应出现 原文件名_aug0.jpg 至 原文件名_aug4.jpg,每张原图对应5张增强图。
  • 视觉差异检查:用图像查看器对比原图与任意增强图,应观察到亮度/对比度变化、轻微平移或旋转(角度≤10°,平移≤8%)。
  • 可复现性检查:固定 seed=42,重复运行两次,生成的图像应完全一致(像素级相同),确保实验可复现。

若需快速预览单张增强效果,可在代码末尾追加以下调试片段:

# 调试:显示单张增强效果(需 matplotlib)
import matplotlib.pyplot as plt
sample = cv2.imread("raw_images/sample.jpg")
sample_rgb = cv2.cvtColor(sample, cv2.COLOR_BGR2RGB)
aug_rgb = cv2.cvtColor(augment_image(sample), cv2.COLOR_BGR2RGB)
plt.figure(figsize=(8,4))
plt.subplot(1,2,1); plt.imshow(sample_rgb); plt.title("Original")
plt.subplot(1,2,2); plt.imshow(aug_rgb); plt.title("Augmented")
plt.show()