import cv2 import numpy as np import os import random import albumentations as A from tqdm import tqdm # --- 1. 用户配置（重点修改！！！）--- # 请根据你的实际路径修改，三个核心目录需区分清楚： # 1. 特征素材来源目录：存放有「待粘贴目标（如no_helmet）」的图片和标签（用于提取可粘贴的目标） SOURCE_FEATURE_IMAGE_DIR = r"E:\NSFW-Detection-YOLO\data\images\val\images" # 有目标的原图 SOURCE_FEATURE_LABEL_DIR = r"E:\NSFW-Detection-YOLO\data\images\val\labels" # 对应原图的标签 # 2. 独立底图目录：存放你要粘贴目标的「空白/背景底图」（底图无需标签） BASE_IMAGE_DIR = r"D:\DataPreHandler\images\valid" # 你的底图文件夹 # 3. 输出目录：保存最终增强后的图片和标签 OUTPUT_IMAGE_DIR = r"D:\DataPreHandler\data\valid\images" OUTPUT_LABEL_DIR = r"D:\DataPreHandler\data\valid\labels" # 数据增强参数 AUGMENTATION_FACTOR = 1 # 每张底图生成的增强图数量（如40张） # --- Copy-Paste 核心配置 --- SMALL_OBJECT_CLASSES_TO_PASTE = [0,1,2,3,4,5,6] # 要粘贴的目标类别ID（如no_helmet是2） PASTE_COUNT_RANGE = (5, 10) # 每张增强图上粘贴的目标数量（随机5-10个） # --- 2. 常规增强流水线（修复Albumentations参数）--- transform_geometric = A.Compose([ A.HorizontalFlip(p=0.5), # 修改1：A.Affine参数：rotate_limit→rotate，cval→pad_val，新增border_mode A.Affine(scale=(0.8, 1.2), shear=(-10, 10), translate_percent=0.1, rotate=30, border_mode=cv2.BORDER_CONSTANT, pad_val=0, p=0.8), A.Perspective(scale=(0.02, 0.05), p=0.4), ], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels'], min_visibility=0.25)) transform_quality = A.Compose([ A.RandomBrightnessContrast(brightness_limit=0.25, contrast_limit=0.25, p=0.8), A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.7), # 修改2：A.GaussNoise参数：var_limit→std_limit（方差转标准差，数值取平方根近似） A.OneOf([A.GaussNoise(std_limit=(3.0, 8.0), p=1.0), A.ISONoise(p=1.0)], p=0.6), A.OneOf([A.Blur(blur_limit=(3, 7), p=1.0), A.MotionBlur(blur_limit=(3, 7), p=1.0)], p=0.5), # 修改3：A.ImageCompression参数：quality_lower/upper→quality_range（合并为元组） A.ImageCompression(quality_range=(70, 95), p=0.3), ], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels'], min_visibility=0.25)) transform_mixed = A.Compose([ A.HorizontalFlip(p=0.5), # 修改4：A.Rotate参数：value→pad_val A.Rotate(limit=15, p=0.5, border_mode=cv2.BORDER_CONSTANT, pad_val=0), A.RandomBrightnessContrast(p=0.6), A.GaussNoise(std_limit=(2.0, 6.0), p=0.4), # 同步修改GaussNoise参数 A.Blur(blur_limit=3, p=0.3), ], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels'], min_visibility=0.25)) base_transforms = [transform_geometric, transform_quality, transform_mixed] # 随机选择增强策略 # --- 3. 核心工具函数 --- def harvest_objects_for_pasting(feature_image_dir, feature_label_dir, target_classes): """ 从「特征素材来源目录」提取目标，创建可粘贴的素材库 :param feature_image_dir: 有目标的图片目录（如含no_helmet的原图） :param feature_label_dir: 对应图片的标签目录 :param target_classes: 要提取的目标类别（如[2]） :return: 素材库 {类别ID: [目标图像1, 目标图像2, ...]} """ print(f"正在从 {feature_image_dir} 提取目标类别 {target_classes}...") asset_library = {cls_id: [] for cls_id in target_classes} # 只读取特征素材目录中的图片文件 feature_image_files = [f for f in os.listdir(feature_image_dir) if f.lower().endswith(('.jpg', '.png', '.jpeg'))] if not feature_image_files: raise FileNotFoundError(f"特征素材目录 {feature_image_dir} 中未找到图片！") for filename in tqdm(feature_image_files, desc="提取目标素材"): label_file = os.path.splitext(filename)[0] + ".txt" label_path = os.path.join(feature_label_dir, label_file) if not os.path.exists(label_path): continue # 跳过无标签的图片 # 读取图片并获取尺寸 img = cv2.imread(os.path.join(feature_image_dir, filename)) if img is None: tqdm.write(f"警告：无法读取图片 {filename}，已跳过") continue img_h, img_w, _ = img.shape # 解析标签，裁剪目标 with open(label_path, 'r') as f: for line in f.readlines(): line = line.strip() if not line: continue parts = line.split() # 修改5：处理标签类别ID为浮点数的情况（如6.0→6）：先转float再转int cls_id = int(float(parts[0])) if cls_id not in target_classes: continue # 只保留目标类别 # YOLO归一化坐标转像素坐标（x1,y1:左上角；x2,y2:右下角） x_center, y_center, box_w, box_h = [float(p) for p in parts[1:]] x1 = int((x_center - box_w / 2) * img_w) y1 = int((y_center - box_h / 2) * img_h) x2 = int((x_center + box_w / 2) * img_w) y2 = int((y_center + box_h / 2) * img_h) # 确保坐标在图片范围内，避免裁剪出错 x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(img_w, x2), min(img_h, y2) # 裁剪目标并加入素材库（排除空图像） if x1 < x2 and y1 < y2: cropped_obj = img[y1:y2, x1:x2] if cropped_obj.size > 0: asset_library[cls_id].append(cropped_obj) # 检查素材库是否为空 total_assets = sum(len(v) for v in asset_library.values()) if total_assets == 0: raise ValueError(f"未从特征素材目录提取到任何目标！请检查类别ID {target_classes} 是否正确") print(f"素材库创建完成！共提取 {total_assets} 个目标（类别：{target_classes}）") return asset_library def paste_objects_to_base(base_image, asset_library): """ 将素材库中的目标粘贴到单张底图上 :param base_image: 输入的底图（cv2读取的BGR图像） :param asset_library: 目标素材库 :return: 粘贴后的图像、对应的YOLO格式标签（bboxes + labels） """ base_h, base_w, _ = base_image.shape pasted_bboxes = [] # 存储粘贴目标的YOLO bbox pasted_labels = [] # 存储粘贴目标的类别ID # 随机确定本次要粘贴的目标数量 num_to_paste = random.randint(*PASTE_COUNT_RANGE) for _ in range(num_to_paste): # 选择要粘贴的目标类别（只从有素材的类别中选） valid_classes = [cls for cls, assets in asset_library.items() if len(assets) > 0] if not valid_classes: break # 极端情况：素材库临时为空（几乎不会发生） # 随机选择一个目标类别和该类别下的一个素材 target_cls = random.choice(valid_classes) target_obj = random.choice(asset_library[target_cls]) obj_h, obj_w, _ = target_obj.shape # 跳过比底图大的目标（避免粘贴后超出边界） if obj_h >= base_h or obj_w >= base_w: continue # 随机选择粘贴位置（左上角坐标，确保目标完全在底图内） paste_x1 = random.randint(0, base_w - obj_w) paste_y1 = random.randint(0, base_h - obj_h) paste_x2 = paste_x1 + obj_w paste_y2 = paste_y1 + obj_h # 直接用Numpy切片粘贴目标（覆盖底图对应区域） base_image[paste_y1:paste_y2, paste_x1:paste_x2] = target_obj # 计算粘贴目标的YOLO归一化坐标（x_center, y_center, w, h） yolo_x_center = (paste_x1 + obj_w / 2) / base_w yolo_y_center = (paste_y1 + obj_h / 2) / base_h yolo_w = obj_w / base_w yolo_h = obj_h / base_h # 将标签加入列表 pasted_bboxes.append([yolo_x_center, yolo_y_center, yolo_w, yolo_h]) pasted_labels.append(target_cls) return base_image, pasted_bboxes, pasted_labels def main(): # 1. 初始化：创建输出目录 os.makedirs(OUTPUT_IMAGE_DIR, exist_ok=True) os.makedirs(OUTPUT_LABEL_DIR, exist_ok=True) # 2. 第一步：创建目标素材库（从特征素材目录提取可粘贴的目标） try: asset_library = harvest_objects_for_pasting( feature_image_dir=SOURCE_FEATURE_IMAGE_DIR, feature_label_dir=SOURCE_FEATURE_LABEL_DIR, target_classes=SMALL_OBJECT_CLASSES_TO_PASTE ) except (FileNotFoundError, ValueError) as e: print(f"错误：{e}") return # 3. 第二步：获取所有底图（只读取图片文件） base_image_files = [f for f in os.listdir(BASE_IMAGE_DIR) if f.lower().endswith(('.jpg', '.png', '.jpeg'))] if not base_image_files: print(f"错误：底图目录 {BASE_IMAGE_DIR} 中未找到任何图片！") return print(f"\n找到 {len(base_image_files)} 张底图，开始生成增强数据（每张底图生成 {AUGMENTATION_FACTOR} 张）") # 4. 主循环：遍历每张底图，生成增强数据 for base_filename in tqdm(base_image_files, desc="处理底图"): base_name, base_ext = os.path.splitext(base_filename) base_image_path = os.path.join(BASE_IMAGE_DIR, base_filename) # 读取底图（若读取失败则跳过） base_image = cv2.imread(base_image_path) if base_image is None: tqdm.write(f"\n警告：无法读取底图 {base_filename}，已跳过") continue # 为当前底图生成 AUGMENTATION_FACTOR 张增强图 for aug_idx in range(AUGMENTATION_FACTOR): # 步骤1：复制底图（避免修改原始底图），并粘贴目标 base_image_copy = base_image.copy() pasted_image, pasted_bboxes, pasted_labels = paste_objects_to_base( base_image=base_image_copy, asset_library=asset_library ) # 步骤2：对粘贴后的图像应用常规增强（Albumentations需要RGB格式） pasted_image_rgb = cv2.cvtColor(pasted_image, cv2.COLOR_BGR2RGB) chosen_transform = random.choice(base_transforms) # 随机选择增强策略 try: # 应用增强（同时处理bbox和label） augmented_result = chosen_transform( image=pasted_image_rgb, bboxes=pasted_bboxes, class_labels=pasted_labels ) final_image_rgb = augmented_result['image'] final_bboxes = augmented_result['bboxes'] final_labels = augmented_result['class_labels'] except Exception as e: tqdm.write(f"\n警告：底图 {base_filename} 增强失败（序号 {aug_idx}）：{str(e)}") continue # 步骤3：保存增强后的图片和标签 # 图片命名格式：底图名_aug_序号.jpg（统一转为jpg格式，避免格式混乱） output_img_name = f"{base_name}_aug_{aug_idx}.jpg" output_img_path = os.path.join(OUTPUT_IMAGE_DIR, output_img_name) # RGB转BGR（cv2保存需要BGR格式） cv2.imwrite(output_img_path, cv2.cvtColor(final_image_rgb, cv2.COLOR_RGB2BGR)) # 标签命名格式：与图片同名.txt（YOLO格式） output_label_name = f"{base_name}_aug_{aug_idx}.txt" output_label_path = os.path.join(OUTPUT_LABEL_DIR, output_label_name) with open(output_label_path, 'w') as f: for bbox, label in zip(final_bboxes, final_labels): x_c, y_c, w, h = bbox # 边界检查：排除增强后可能超出0-1范围的bbox（避免训练报错） if 0 <= x_c <= 1 and 0 <= y_c <= 1 and 0 <= w <= 1 and 0 <= h <= 1: f.write(f"{label} {x_c:.6f} {y_c:.6f} {w:.6f} {h:.6f}\n") # 5. 完成提示 total_generated = len(base_image_files) * AUGMENTATION_FACTOR print(f"\n✅ 数据增强全部完成！") print(f"📊 生成数据统计：") print(f" - 底图数量：{len(base_image_files)} 张") print(f" - 每张底图增强次数：{AUGMENTATION_FACTOR} 次") print(f" - 总生成图片/标签：{total_generated} 组") print(f" - 输出路径：") print(f" 图片 → {OUTPUT_IMAGE_DIR}") print(f" 标签 → {OUTPUT_LABEL_DIR}") if __name__ == "__main__": # 运行前务必确认： # 1. SOURCE_FEATURE_IMAGE_DIR/SOURCE_FEATURE_LABEL_DIR 是「有目标的素材目录」 # 2. BASE_IMAGE_DIR 是你的「空白底图目录」 # 3. SMALL_OBJECT_CLASSES_TO_PASTE 是要粘贴的目标类别ID（如no_helmet=2） main()