代码收藏家 YOLO 数据增强 Python 脚本（可选次数，无限随机增强）- 一键执行搞定，自动化提升训练集质量 | 幽络源
前言
往往在准备需要训练一个模型的时候，很多人苦于找不到合适的数据集，自己标注又耗时耗力，而数据增强正好解决了这个问题，因此对于数据增强这个概念是非常有必要的，本文将提供一个数据增强脚本，你无需理解代码，只需懂得如何使用即可达到你要的效果。
背景
近期我在一直寻找冲沟相关的Box标注数据集，结果是收费的与不收费的都没有找到，唯一找到的和冲沟相关还是一个图像分类的数据集，不是Box标注的对象识别的，如下图
以上图片其实很符合冲沟相关的特征与缺陷，只是这些图没有被标注，那就只有自己标注了，标注完后，前前后后加起来也才125张，这点数据量相对于这种复杂的地貌和缺陷是远远不够的，因此我突然想起一个数据增强的概率，或许会有用（我是外行）。
图片数据增强脚本
先直接上源码，然后再说使用方式
import os
import cv2
import albumentations as A
from tqdm import tqdm
# 数据增强脚本
def augment_images_and_labels(img_dir, label_dir, output_img_dir, output_label_dir, val_img_dir, val_label_dir, augment_times=3):
"""
对YOLO数据进行安全的数据增强。
:param img_dir: 原始图片的目录路径
:param label_dir: YOLO标签目录路径
:param output_img_dir: 增强后的图片保存目录
:param output_label_dir: 增强后的标签保存目录
:param val_img_dir: 第一个增强图片存放的目录，用于验证集
:param val_label_dir: 第一个增强标签存放的目录，用于验证集
:param augment_times: 每张图片的增强次数
"""
# 确保输出目录存在
os.makedirs(output_img_dir, exist_ok=True)
os.makedirs(output_label_dir, exist_ok=True)
os.makedirs(val_img_dir, exist_ok=True)
os.makedirs(val_label_dir, exist_ok=True)
# 遍历图片和标签
for img_file in tqdm(os.listdir(img_dir)):
if not img_file.endswith(('.jpg', '.png', '.jpeg')):
continue
# 获取文件路径
img_path = os.path.join(img_dir, img_file)
label_path = os.path.join(label_dir, os.path.splitext(img_file)[0] + '.txt')
# 读取图像
image = cv2.imread(img_path)
if image is None:
print(f"无法读取图像: {img_path}")
continue
height, width = image.shape[:2]
# 读取YOLO标签
bboxes = []
class_labels = []
if os.path.exists(label_path):
with open(label_path, 'r') as f:
for line in f.readlines():
parts = line.strip().split()
class_id = int(parts[0])
x_center, y_center, w, h = map(float, parts[1:])
bboxes.append([x_center, y_center, w, h])
class_labels.append(class_id)
# 原始标签检查
if not bboxes:
print(f"标签为空: {label_path}")
continue
# 进行数据增强
for i in range(augment_times):
# 动态调整裁剪区域大小（最大不能大于图像尺寸）
min_crop_size = min(height, width)
crop_height = min(500, min_crop_size)
crop_width = min(500, min_crop_size)
# 更新增强方法中的裁剪尺寸
augmentations = A.Compose([
A.HorizontalFlip(p=0.5),
A.RandomBrightnessContrast(p=0.5),
A.Rotate(limit=10, p=0.5, border_mode=cv2.BORDER_CONSTANT),
A.GaussianBlur(p=0.2),
A.GaussNoise(p=0.2),
A.Resize(width=640, height=640, p=0.5),
A.RandomCrop(width=crop_width, height=crop_height, p=0.5),
A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.5),
A.ElasticTransform(p=0.2),
A.RandomScale(p=0.2),
], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))
augmented = augmentations(image=image, bboxes=bboxes, class_labels=class_labels)
aug_image = augmented['image']
aug_bboxes = augmented['bboxes']
aug_labels = augmented['class_labels']
# 如果是第一个增强，保存到验证集目录
if i == 0:
output_img_path = os.path.join(val_img_dir, f"{os.path.splitext(img_file)[0]}_aug_0.jpg")
output_label_path = os.path.join(val_label_dir, f"{os.path.splitext(img_file)[0]}_aug_0.txt")
else:
# 否则，保存到训练集目录
output_img_path = os.path.join(output_img_dir, f"{os.path.splitext(img_file)[0]}_aug_{i}.jpg")
output_label_path = os.path.join(output_label_dir, f"{os.path.splitext(img_file)[0]}_aug_{i}.txt")
# 保存增强后的图片
cv2.imwrite(output_img_path, aug_image)
# 保存增强后的标签
with open(output_label_path, 'w') as f:
for bbox, cls in zip(aug_bboxes, aug_labels):
x_center, y_center, w, h = bbox
f.write(f"{int(cls)} {x_center} {y_center} {w} {h}\n")
print("数据增强完成！")
if __name__ == "__main__":
# 输入目录路径，必须是全英文路径
img_dir = r"D:\A01PythonProjects3123\labelImg-master\gullyYolo\images" # 原始图片路径
label_dir = r"D:\A01PythonProjects3123\labelImg-master\gullyYolo\labels" # YOLO标签路径
# 输出目录路径
output_img_dir = "handleImages" # 保存增强后图片路径
output_label_dir = "handleLabels" # 保存增强后标签路径
# 验证集目录路径
val_img_dir = "val_images" # 保存增强后验证集图片路径
val_label_dir = "val_labels" # 保存增强后验证集标签路径
# 执行数据增强
augment_images_and_labels(img_dir, label_dir, output_img_dir, output_label_dir, val_img_dir, val_label_dir, augment_times=8)使用方式
对于以上的源码，你唯一需要修改的便是 main方法中的img_dir与label_dir的路径，分别为你需要增强的图片目录路径和图片对应的标签目录路径。输出目录和验证集目录不用修改，然后你可以通过修改 augment_images_and_labels函数中的augment_times的次数，这个次数决定了增强的次数，每次增强都是不一样的效果。
确定好这些即可执行脚本完成数据的增强，如图，执行后输出了四个目录，红框中的两个目录作为训练集使用，黄框中的两个目录作为验证集使用
我这里的次数设定为了8，最终handleImages的目录中的图片如下
虽然次数为8，但是这作为训练集的每张图只具备了7张，因为还有一张作为了验证集，这样能保证每一张图都能得到好的训练+验证，并且我的图片由原来的125张变为了现在的984张
本脚本优势
增强操作多样
albumentations库，支持多种常见的图像增强操作，如：
HorizontalFlip)RandomBrightnessContrast)Rotate)GaussianBlur)GaussNoise)RandomCrop)HueSaturationValue)ElasticTransform)RandomScale)验证集与训练集区分
i == 0）被保存到专门的验证集目录（val_img_dir和val_label_dir），其余增强样本被保存到训练集目录（output_img_dir和output_label_dir）。这种做法保证了验证集与训练集数据的区分，避免验证集样本被重复用于训练。动态裁剪大小调整
可定制的增强次数
augment_times参数控制，默认值为3。用户可以根据需求调整这个值，以控制最终生成的增强数据量。数据增强的好处
1. 增加数据多样性
2. 提升模型的泛化能力
3. 提高模型对噪声的鲁棒性
4. 扩展训练数据集
5. 应对标签稀缺问题
6. 优化数据质量
7. 应对不平衡数据问题
8. 模拟现实场景
9. 加速训练过程
总结
以上是幽络源对数据增强的应用与理解，个人认为做数据集，数据增强是非常有必要的
作者：幽络源小助理
