将图片进行局部直方图均衡化

这段代码可以使强光、反光的图片亮度变暗，可以使强光、反光出现的误报漏报，进行局部直方图均衡化后更好的进行训练。

import numpy as np
import cv2
import os

# 定义一个函数augment_hsv，用于增强图像的HSV颜色空间。
def augment_hsv(img, hgain=0.015, sgain=0.7, vgain=0.4):
    # 生成随机的HSV增益系数。这些系数用于调整图像的色调、饱和度和亮度。
    r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1
    # 将图像从BGR颜色空间转换为HSV颜色空间。
    hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
    # 获取图像的数据类型。
    dtype = img.dtype
    # 创建一个从0到255的数组，数据类型为uint8（无符号8位整数）。
    x = np.arange(0, 256, dtype=np.int16)
    # 根据随机增益系数创建一个查找表（LUT），用于调整色调。
    lut_hue = ((x * r[0]) % 180).astype(dtype)
    # 创建一个查找表，用于调整饱和度。
    lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
    # 创建一个查找表，用于调整亮度。
    lut_val = np.clip(x * r[2], 0, 255).astype(dtype)
    # 使用查找表对色调、饱和度和亮度进行调整。
    img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))).astype(dtype)
    # 将调整后的HSV图像转换回BGR颜色空间。
    cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img)
    # 返回处理后的图像。注意这里不需要返回值，因为已经直接修改了传入的图像变量。
    return img_hsv

# 定义一个函数data_augmentation，用于对图像进行直方图均衡化和高斯模糊。
def data_augmentation(image_path):
    # 从指定路径读取图像。
    image = cv2.imread(image_path)
    # 将图像从BGR颜色空间转换为YUV颜色空间。
    img_yuv = cv2.cvtColor(image, cv2.COLOR_BGR2YUV)
    # 对图像的亮度通道进行直方图均衡化，以增强对比度。
    img_yuv[:, :, 0] = cv2.equalizeHist(img_yuv[:, :, 0])
    # 对图像应用高斯模糊，以减少噪声。
    img_yuv = cv2.GaussianBlur(img_yuv, (3, 3), 0)
    # 将图像从YUV颜色空间转换回BGR颜色空间。
    image = cv2.cvtColor(img_yuv, cv2.COLOR_YUV2BGR)
    return image

# 定义一个主函数，用于执行整个程序的主要逻辑。
def main():
    # 设置输入图片的文件夹路径。
    input_folder="/home/zhang/yolov5_obb_master/old_bottles/images1"
    # 设置输出图片的文件夹路径。如果路径不存在，则创建它。
    output_folder=r"/home/zhang/yolov5_obb_master/old_bottles/images"
    # 检查输入文件夹是否存在，如果不存在则创建。
    if not os.path.exists(input_folder):
        os.makedirs(output_folder)
    # 遍历输入文件夹中的所有文件。
    for filename in os.listdir(input_folder):
        # 只处理jpg和png格式的图片文件。
        if filename.endswith("jpg") or filename.endswith("png"):
            # 调用data_augmentation函数对图像进行处理。
            augmented_image=data_augmentation(os.path.join(input_folder,filename))
            # 将处理后的图像保存到输出文件夹中，文件名保持不变，只是格式会变为png（因为cv2.imwrite默认保存为png格式）。
            cv2.imwrite(os.path.join(output_folder,filename[:-4]+filename[-4:]),
                        augmented_image)
            # 打印一条消息，告知用户该图像已被处理并保存。
            print(f"Processed {filename} and saved the augmented version as {filename[:-4]+filename[-4:]}")

# 如果这个脚本是直接运行的，而不是被其他脚本导入的，那么就执行main函数。
if __name__ == '__main__':
    main()