接上个博客：
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)