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)
运行方式
- 准备一个存放原始图像的目录(如
raw_images/),放入若干.jpg或.png前视驾驶图像。 - 将上述代码保存为
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()