yolo训练策略--使用 Python 和 OpenCV 进行图像亮度增强与批量文件复制之(图像增强是按梯度变化优化)

接上个博客:

复制代码
https://blog.csdn.net/weixin_43269994/article/details/141753412

优化如下函数:

复制代码
def augment_and_copy_files(base_folder, image_filename, num_augmentations=2, vgain_range=(1, 1.5), process_labels=True, process_annotations=True):
    base_filename, image_ext = os.path.splitext(image_filename)

    # 构建原始文件路径
    file_paths = {
        "images": os.path.join(base_folder, "images", image_filename),
    }

    if process_annotations:
        file_paths["annotations"] = os.path.join(base_folder, "annotations", f"{base_filename}.xml")
    if process_labels:
        file_paths["labels"] = os.path.join(base_folder, "labels", f"{base_filename}.txt")

    # 创建输出文件夹
    output_folders = create_output_folders(base_folder)

    # 复制原始图像
    copy_file(file_paths["images"], output_folders["images"], "", preserve_ext=True)

    if process_annotations:
        copy_file(file_paths["annotations"], output_folders["annotations"], "", preserve_ext=True)
    if process_labels:
        copy_file(file_paths["labels"], output_folders["labels"], "", preserve_ext=True)

    # 生成按梯度变化的增益值
    vgain_start, vgain_end = vgain_range
    vgain_step = (vgain_end - vgain_start) / num_augmentations

    for i in range(1, num_augmentations + 1):
        vgain = vgain_start + i * vgain_step
        brightened_img = adjust_brightness(cv2.imread(file_paths["images"]), vgain)

        filename_suffix = f"_enhanced_{i}"
        output_image_path = copy_file(file_paths["images"], output_folders["images"], filename_suffix, preserve_ext=True)
        cv2.imwrite(output_image_path, brightened_img)
        print(f"Saved: {output_image_path}")

        if process_annotations:
            copy_file(file_paths["annotations"], output_folders["annotations"], filename_suffix, preserve_ext=True)
            print(f"Copied annotations: {output_image_path}")

        if process_labels:
            copy_file(file_paths["labels"], output_folders["labels"], filename_suffix, preserve_ext=True)
            print(f"Copied labels: {output_image_path}")

    print(f"All unique images and their annotations for {image_filename} have been enhanced and saved!")

这个函数 augment_and_copy_files 的目的是处理和增强图像,并将处理后的图像及其相关的注释和标签文件复制到指定的输出文件夹中。具体来说,它对图像进行亮度调整,并生成多个增强版本,同时可选择处理和复制对应的注释和标签文件。以下是详细解释:

  • base_folder: 原始数据的基路径。它包含了 images、annotations 和 labels 文件夹。
  • image_filename: 要处理的图像文件名。
  • num_augmentations: 生成的增强图像数量。
  • vgain_range: 亮度增益的范围,包含两个值,起始增益和结束增益。
  • process_labels: 布尔值,指示是否处理标签文件。
  • process_annotations: 布尔值,指示是否处理注释文件。

总体代码:

复制代码
import cv2
import numpy as np
import os
import shutil


def adjust_brightness(im, vgain):
    hsv = cv2.cvtColor(im, cv2.COLOR_BGR2HSV)
    hue, sat, val = cv2.split(hsv)
    val = np.clip(val * vgain, 0, 255).astype(np.uint8)
    enhanced_hsv = cv2.merge((hue, sat, val))
    brightened_img = cv2.cvtColor(enhanced_hsv, cv2.COLOR_HSV2BGR)
    return brightened_img


def create_output_folders(base_folder):
    new_base_folder = os.path.join(os.path.dirname(base_folder), "augmented_data")
    output_folders = {
        "images": os.path.join(new_base_folder, "images"),
        "annotations": os.path.join(new_base_folder, "annotations"),
        "labels": os.path.join(new_base_folder, "labels")
    }
    for folder in output_folders.values():
        os.makedirs(folder, exist_ok=True)
    return output_folders


def copy_file(src_path, dst_folder, filename_suffix, preserve_ext=True):
    base_filename, ext = os.path.splitext(os.path.basename(src_path))
    if preserve_ext:
        new_filename = f"{base_filename}{filename_suffix}{ext}"
    else:
        new_filename = f"{base_filename}{filename_suffix}"
    dst_path = os.path.join(dst_folder, new_filename)
    shutil.copy(src_path, dst_path)
    return dst_path


def augment_and_copy_files(base_folder, image_filename, num_augmentations=2, vgain_range=(1, 1.5), process_labels=True, process_annotations=True):
    base_filename, image_ext = os.path.splitext(image_filename)

    # 构建原始文件路径
    file_paths = {
        "images": os.path.join(base_folder, "images", image_filename),
    }

    if process_annotations:
        file_paths["annotations"] = os.path.join(base_folder, "annotations", f"{base_filename}.xml")
    if process_labels:
        file_paths["labels"] = os.path.join(base_folder, "labels", f"{base_filename}.txt")

    # 创建输出文件夹
    output_folders = create_output_folders(base_folder)

    # 复制原始图像
    copy_file(file_paths["images"], output_folders["images"], "", preserve_ext=True)

    if process_annotations:
        copy_file(file_paths["annotations"], output_folders["annotations"], "", preserve_ext=True)
    if process_labels:
        copy_file(file_paths["labels"], output_folders["labels"], "", preserve_ext=True)

    # 生成按梯度变化的增益值
    vgain_start, vgain_end = vgain_range
    vgain_step = (vgain_end - vgain_start) / num_augmentations

    for i in range(1, num_augmentations + 1):
        vgain = vgain_start + i * vgain_step
        brightened_img = adjust_brightness(cv2.imread(file_paths["images"]), vgain)

        filename_suffix = f"_enhanced_{i}"
        output_image_path = copy_file(file_paths["images"], output_folders["images"], filename_suffix, preserve_ext=True)
        cv2.imwrite(output_image_path, brightened_img)
        print(f"Saved: {output_image_path}")

        if process_annotations:
            copy_file(file_paths["annotations"], output_folders["annotations"], filename_suffix, preserve_ext=True)
            print(f"Copied annotations: {output_image_path}")

        if process_labels:
            copy_file(file_paths["labels"], output_folders["labels"], filename_suffix, preserve_ext=True)
            print(f"Copied labels: {output_image_path}")

    print(f"All unique images and their annotations for {image_filename} have been enhanced and saved!")


def process_all_images_in_folder(base_folder, num_augmentations=2, vgain_range=(1, 1.5), process_labels=True, process_annotations=True):
    images_folder = os.path.join(base_folder, "images")
    for image_filename in os.listdir(images_folder):
        if image_filename.lower().endswith(('.bmp', '.jpg', '.jpeg', '.png')):
            augment_and_copy_files(base_folder, image_filename, num_augmentations, vgain_range, process_labels, process_annotations)


# 使用示例
base_folder = r"C:\Users\linds\Desktop\fsdownload\upgrade_algo_so\data_res_2024_08_31_16_38\train"
process_all_images_in_folder(base_folder, num_augmentations=10, vgain_range=(1, 3), process_labels=True, process_annotations=False)
相关推荐
mit6.824几秒前
[GazeTracking] 摄像头交互与显示 | OpenCV
人工智能·opencv·交互
工业互联网专业15 分钟前
基于大数据的学习资源推送系统的设计与实现 _django
vue.js·python·django·毕业设计·源码·课程设计·学习资源推送系统
木子杳衫3 小时前
【软件开发】管理类系统
python·web开发
程序员小远6 小时前
银行测试:第三方支付平台业务流,功能/性能/安全测试方法
自动化测试·软件测试·python·功能测试·测试工具·性能测试·安全性测试
星期天要睡觉7 小时前
计算机视觉(opencv)——基于 OpenCV DNN 的实时人脸检测 + 年龄与性别识别
opencv·计算机视觉·dnn
猫头虎8 小时前
如何查看局域网内IP冲突问题?如何查看局域网IP环绕问题?arp -a命令如何使用?
网络·python·网络协议·tcp/ip·开源·pandas·pip
沿着路走到底8 小时前
python 基础
开发语言·python
烛阴10 小时前
武装你的Python“工具箱”:盘点10个你必须熟练掌握的核心方法
前端·python
杨枝甘露小码11 小时前
Python学习之基础篇
开发语言·python
我是华为OD~HR~栗栗呀11 小时前
23届考研-Java面经(华为OD)
java·c++·python·华为od·华为·面试