Python—4.yolo的图片及带xtx标签实现数据增强【符完整代码】

这个代码主要是解析训练yolo的xtx标签实现数据增强中能通过改变色调：包括随机亮度、对比度和饱和度的调整。

文件夹路径：input_image_dir、input_label_dir和output_image_dir、output_label_dir分别设置为指定的图像、标签和输出指定的图像、标签路径。
增强整个文件夹：augment_images函数遍历输入文件夹中的所有图像文件，对每个图像及其对应的标签进行数据增强，并将结果保存到输出文件夹中。
创建目录：使用os.makedirs确保输出目录存在。

代码详解
主要是修改input_image_dir、input_label_dir和output_image_dir、output_label_dir这几个地方即可，话不多说 ，咱们直接附完整代码：

import cv2
import random
import os

def load_labels(label_path):
    """加载YOLO标签文件"""
    with open(label_path, 'r') as file:
        lines = file.readlines()
    return [line.strip() for line in lines]

def save_labels(labels, label_path):
    """保存YOLO标签文件"""
    with open(label_path, 'w') as file:
        for line in labels:
            file.write(line + '\n')

def augment_image(image_path, label_path, output_image_dir, output_label_dir, augmentation_index, flagcount, json_names):
    """对图像进行数据增强并更新标签"""
    # 加载图像和标签
    image = cv2.imread(image_path)
    labels = load_labels(label_path)
    
    # 2. 改变色调
    # 随机亮度调整
    brightness = random.uniform(0.5, 1.5)
    image = cv2.convertScaleAbs(image, alpha=1, beta=int(brightness * 100 - 100))
    
    # 随机对比度调整
    contrast = random.uniform(0.5, 1.5)
    image = cv2.convertScaleAbs(image, alpha=contrast, beta=0)
    
    # 随机饱和度调整
    saturation = random.uniform(0.5, 1.5)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    image[:, :, 1] = image[:, :, 1] * saturation
    image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
    
    # 保存增强后的图像和标签
    base_name = os.path.basename(image_path)
    new_image_path = os.path.join(output_image_dir, f"{os.path.splitext(base_name)[0]}_{augmentation_index}.jpg")
    new_label_path = os.path.join(output_label_dir, f"{os.path.splitext(base_name)[0]}_{augmentation_index}.txt")
    os.makedirs(os.path.dirname(new_image_path), exist_ok=True)
    os.makedirs(os.path.dirname(new_label_path), exist_ok=True)
    cv2.imwrite(new_image_path, image)
    save_labels(labels, new_label_path)
    flagcount += 1  # 更新已处理文件计数
    return flagcount

def augment_images(input_image_dir, input_label_dir, output_image_dir, output_label_dir):
    """对整个文件夹中的图像进行数据增强"""
    json_names = os.listdir(input_image_dir)  # 获取所有文件名
    flagcount = 0  # 初始化已处理文件计数
    for filename in json_names:
        if filename.endswith(".jpg") or filename.endswith(".png"):
            image_path = os.path.join(input_image_dir, filename)
            label_path = os.path.join(input_label_dir, os.path.splitext(filename)[0] + '.txt')
            for i in range(3):  # 增强三次
                flagcount = augment_image(image_path, label_path, output_image_dir, output_label_dir, i+1, flagcount, json_names)
    print("还剩下{}个文件未转化".format(len(json_names) - flagcount))
    print('转化全部完毕')

# 使用示例
input_image_dir = r'D:\Users\19114\anaconda3\envs\yolo\Lib\site-packages\ultralytics\SK-MPSO-BP\tata5\images'
input_label_dir = r'D:\Users\19114\anaconda3\envs\yolo\Lib\site-packages\ultralytics\SK-MPSO-BP\tata5\labels'
output_image_dir = r'D:\Users\19114\anaconda3\envs\yolo\Lib\site-packages\ultralytics\SK-MPSO-BP\tata5\new_images'
output_label_dir = r'D:\Users\19114\anaconda3\envs\yolo\Lib\site-packages\ultralytics\SK-MPSO-BP\tata5\new_labels'
augment_images(input_image_dir, input_label_dir, output_image_dir, output_label_dir)

说明
运行这段代码后，每张图片将进行三次色调调整的数据增强，增强后的图片和标签将分别保存到指定的输出文件夹中，并且文件名后会加上_1、_2、_3等后缀。同时，代码会在处理完毕后打印剩余未处理的文件数量和转化完成的消息。

增强次数：在augment_images函数中，通过一个循环for i in range(3)来确保每张图片至少增强三次。
文件命名：在augment_image函数中，增强后的图片和标签的文件名后会加上后缀_{augmentation_index}，其中augmentation_index从1开始，以区分不同的增强版本。
文件处理计数：flagcount用于跟踪已处理的文件数量。
剩余文件打印：在所有文件处理完毕后，打印剩余未处理的文件数量，并输出转化完成的消息。