yolo训练策略--使用 Python 和 OpenCV 进行图像亮度增强与批量文件复制

简介

在计算机视觉和深度学习项目中,数据增强是一种常用的技术,通过对原始图像进行多种变换,可以增加数据集的多样性,从而提高模型的泛化能力。本文将介绍如何使用 Python 和 OpenCV 实现图像的亮度增强,并将增强后的图像与对应的注释文件批量复制到新目录中。

项目背景

假设你有一个数据集,包含若干图像及其对应的 XML 注释文件和标签文件。在模型训练前,你希望对这些图像进行亮度增强,并生成新的图像及其对应的注释文件和标签文件。本教程将指导你如何编写一个 Python 脚本,实现此功能。

train目录如下:

生成的augmented_data如下:

代码实现

1. 图像亮度调整函数

首先,我们需要编写一个函数,来调整图像的亮度。此处我们使用 HSV 色彩空间的 V(亮度)通道进行调整。

import cv2
import numpy as np

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

2. 创建输出目录

在进行文件操作前,我们需要为增强后的文件创建一个新的输出目录。

import os

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

3. 文件复制函数

为了复制原始图像和对应的注释文件,我们编写一个通用的文件复制函数。该函数可以根据需要在文件名后添加后缀。

import shutil

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

4. 图像增强与文件复制

该函数实现了图像的亮度增强,同时将增强后的图像和对应的注释文件保存到新的目录中。

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

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

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

    # 复制原始文件
    for key in file_paths:
        copy_file(file_paths[key], output_folders[key], "", preserve_ext=True)

    # 确保增强结果不重复
    unique_vgains = set()

    while len(unique_vgains) < num_augmentations:
        vgain = np.random.uniform(*vgain_range)
        if vgain not in unique_vgains:
            unique_vgains.add(vgain)
            brightened_img = adjust_brightness(cv2.imread(file_paths["images"]), vgain)

            for key in file_paths:
                filename_suffix = f"_enhanced_{len(unique_vgains)}"
                output_path = copy_file(file_paths[key], output_folders[key], filename_suffix, preserve_ext=True)
                if key == "images":
                    cv2.imwrite(output_path, brightened_img)
                    print(f"Saved: {output_path}")
                else:
                    print(f"Copied {key}: {output_path}")

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

5. 处理整个目录

最后,我们编写一个函数,用于处理指定目录中的所有图像文件,并对每张图像进行增强。

def process_all_images_in_folder(base_folder, num_augmentations=2, vgain_range=(1, 1.5)):
    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)

6. 运行脚本

你可以通过以下代码来运行整个图像增强与文件复制过程:

# 使用示例
base_folder = r"C:\Users\linds\Desktop\fsdownload\upgrade_algo_so\data_res_2024_08_31_10_29\train"
process_all_images_in_folder(base_folder)

7.整体代码

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)):
    base_filename, image_ext = os.path.splitext(image_filename)

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

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

    # 复制原始文件
    for key in file_paths:
        copy_file(file_paths[key], output_folders[key], "", preserve_ext=True)

    # 确保增强结果不重复
    unique_vgains = set()

    while len(unique_vgains) < num_augmentations:
        vgain = np.random.uniform(*vgain_range)
        if vgain not in unique_vgains:
            unique_vgains.add(vgain)
            brightened_img = adjust_brightness(cv2.imread(file_paths["images"]), vgain)

            for key in file_paths:
                filename_suffix = f"_enhanced_{len(unique_vgains)}"
                output_path = copy_file(file_paths[key], output_folders[key], filename_suffix, preserve_ext=True)
                if key == "images":
                    cv2.imwrite(output_path, brightened_img)
                    print(f"Saved: {output_path}")
                else:
                    print(f"Copied {key}: {output_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)):
    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)


# 使用示例
base_folder = r"C:\Users\linds\Desktop\fsdownload\upgrade_algo_so\data_res_2024_08_31_10_29\train"
process_all_images_in_folder(base_folder)
相关推荐
龙哥说跨境15 分钟前
如何利用指纹浏览器爬虫绕过Cloudflare的防护?
服务器·网络·python·网络爬虫
小白学大数据31 分钟前
正则表达式在Kotlin中的应用:提取图片链接
开发语言·python·selenium·正则表达式·kotlin
flashman91133 分钟前
python在word中插入图片
python·microsoft·自动化·word
菜鸟的人工智能之路36 分钟前
桑基图在医学数据分析中的更复杂应用示例
python·数据分析·健康医疗
懒大王爱吃狼2 小时前
Python教程:python枚举类定义和使用
开发语言·前端·javascript·python·python基础·python编程·python书籍
秃头佛爷3 小时前
Python学习大纲总结及注意事项
开发语言·python·学习
深度学习lover4 小时前
<项目代码>YOLOv8 苹果腐烂识别<目标检测>
人工智能·python·yolo·目标检测·计算机视觉·苹果腐烂识别
API快乐传递者5 小时前
淘宝反爬虫机制的主要手段有哪些?
爬虫·python
阡之尘埃7 小时前
Python数据分析案例61——信贷风控评分卡模型(A卡)(scorecardpy 全面解析)
人工智能·python·机器学习·数据分析·智能风控·信贷风控
Eric.Lee20219 小时前
yolo v5 开源项目
人工智能·yolo·目标检测·计算机视觉