import cv2 import numpy as np import os import random from collections import defaultdict from tqdm import tqdm # --- 1. 核心配置 --- # 特征素材来源目录（包含完整图片和对应标签） # SOURCE_FEATURE_IMAGE_DIR = r"D:\DataPreHandler\yuanshi_data\images\train\images" # SOURCE_FEATURE_LABEL_DIR = r"D:\DataPreHandler\yuanshi_data\images\train\labels" # SOURCE_FEATURE_IMAGE_DIR = r"D:\DataPreHandler\yuanshi_data\images\val\images" # SOURCE_FEATURE_LABEL_DIR = r"D:\DataPreHandler\yuanshi_data\images\val\labels" SOURCE_FEATURE_IMAGE_DIR = r"D:\DataPreHandler\yuanshi_data\images\test\images" SOURCE_FEATURE_LABEL_DIR = r"D:\DataPreHandler\yuanshi_data\images\test\labels" # 底图目录 # BASE_IMAGE_DIR = r"D:\DataPreHandler\images\dituchoqu\train" BASE_IMAGE_DIR = r"D:\DataPreHandler\images\dituchoqu_huashen\train" # BASE_IMAGE_DIR = r"D:\DataPreHandler\images\test" # 输出目录 OUTPUT_IMAGE_DIR = r"D:\DataPreHandler\data\train3\images" OUTPUT_LABEL_DIR = r"D:\DataPreHandler\data\train3\labels" # OUTPUT_IMAGE_DIR = r"D:\DataPreHandler\data\val3\images" # OUTPUT_LABEL_DIR = r"D:\DataPreHandler\data\val3\labels" # OUTPUT_IMAGE_DIR = r"D:\DataPreHandler\data\test\da\images" # OUTPUT_LABEL_DIR = r"D:\DataPreHandler\data\test\da\labels" # 底图数量 SELECT_BASE_IMAGE_COUNT = 0 AUGMENTATION_FACTOR = 1 TARGET_CLASSES = [0, 1, 2, 3, 4, 5, 6] PASTE_COUNT_RANGE = (2, 3) MIN_SCALED_SIZE = 50 MAX_OVERLAP_IOU = 0.0 # 仅限制“跨特征图”无重叠 MAX_RETRY_COUNT = 300 SCALE_FACTOR_RANGE = (0.3, 0.9) # --- 2. 工具函数 ---（核心修改：删除内部重叠检测） def calculate_iou(box1, box2): """计算两个边界框的交并比（像素坐标）——保持不变""" box1 = [int(round(x)) for x in box1] box2 = [int(round(x)) for x in box2] inter_x1 = max(box1[0], box2[0]) inter_y1 = max(box1[1], box2[1]) inter_x2 = min(box1[2], box2[2]) inter_y2 = min(box1[3], box2[3]) inter_width = max(0, inter_x2 - inter_x1) inter_height = max(0, inter_y2 - inter_y1) inter_area = inter_width * inter_height box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1]) box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1]) union_area = box1_area + box2_area - inter_area return inter_area / union_area if union_area > 0 else 0.0 def collect_feature_images(feature_image_dir, feature_label_dir, target_classes): """收集特征图片及其标签——核心修改：删除内部重叠过滤""" print(f"从 {feature_image_dir} 收集特征图片及标签，目标类别 {target_classes}...") feature_list = [] # 元素格式：(图片路径, 标签列表, 原始宽, 原始高) 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"未找到特征图片！") for filename in tqdm(feature_image_files, desc="收集特征图片"): img_path = os.path.join(feature_image_dir, filename) 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(img_path) if img is None: tqdm.write(f"警告：无法读取图片 {filename}，已跳过") continue orig_h, orig_w = img.shape[:2] # 仅过滤缩小后可能过小的特征图（保留内部重叠的图） min_required_orig_size = MIN_SCALED_SIZE / max(SCALE_FACTOR_RANGE) if orig_w < min_required_orig_size or orig_h < min_required_orig_size: tqdm.write(f"警告：特征图片 {filename} 原始尺寸过小，已跳过") continue # 读取并过滤目标类别（不再检查内部重叠） labels = [] with open(label_path, 'r') as f: for line in f.readlines(): line = line.strip() if not line: continue parts = line.split() cls_id = int(float(parts[0])) if cls_id not in target_classes: continue xc, yc, w, h = [float(p) for p in parts[1:]] if 0 <= xc <= 1 and 0 <= yc <= 1 and 0 < w <= 1 and 0 < h <= 1: labels.append((cls_id, xc, yc, w, h)) # 只要有有效标签就保留（无论内部是否重叠） if labels: feature_list.append((img_path, labels, orig_w, orig_h)) if not feature_list: raise ValueError(f"未收集到有效的特征图片及标签！") print(f"特征图片收集完成：共 {len(feature_list)} 张有效特征图片（允许内部目标重叠）") return feature_list def paste_feature_images_to_base(base_image, feature_list): """将特征图粘贴到底图——保留“跨特征图无重叠”逻辑（核心）""" base_h, base_w = base_image.shape[:2] pasted_labels = [] # 最终输出的标签 pasted_target_boxes = [] # 已粘贴的目标框（像素坐标） pasted_feature_regions = [] # 已粘贴的特征图整体区域（防跨图重叠） pasted_feature_count = 0 # 成功粘贴的特征图数量 skipped_small = 0 # 因尺寸过小跳过的次数 skipped_region_overlap = 0 # 因特征图整体区域重叠跳过的次数 skipped_target_overlap = 0 # 因跨图目标重叠跳过的次数 retry_count = 0 # 重试次数 target_paste_count = random.randint(*PASTE_COUNT_RANGE) print(f" 计划粘贴 {target_paste_count} 张特征图...") while pasted_feature_count < target_paste_count and retry_count < MAX_RETRY_COUNT: retry_count += 1 # 1. 随机选择一张特征图（允许内部重叠） feature_img_path, feature_labels, orig_w, orig_h = random.choice(feature_list) # 2. 随机缩放特征图，确保尺寸符合要求 scale_factor = random.uniform(*SCALE_FACTOR_RANGE) scaled_w = int(round(orig_w * scale_factor)) scaled_h = int(round(orig_h * scale_factor)) if scaled_w < MIN_SCALED_SIZE or scaled_h < MIN_SCALED_SIZE: skipped_small += 1 continue if scaled_w >= base_w or scaled_h >= base_h: skipped_small += 1 continue # 3. 随机生成粘贴位置（确保特征图完全在底图内） paste_x1 = random.randint(0, base_w - scaled_w) paste_y1 = random.randint(0, base_h - scaled_h) paste_x2 = paste_x1 + scaled_w paste_y2 = paste_y1 + scaled_h current_feature_region = (paste_x1, paste_y1, paste_x2, paste_y2) # 4. 关键：检测当前特征图整体区域与已粘贴区域是否重叠（防跨图重叠） region_overlap = False for existing_region in pasted_feature_regions: if calculate_iou(current_feature_region, existing_region) > MAX_OVERLAP_IOU: region_overlap = True skipped_region_overlap += 1 break if region_overlap: continue # 5. 计算当前特征图目标在底图上的像素坐标（保留原始内部重叠） temp_target_boxes = [] valid = True for (cls_id, xc, yc, w, h) in feature_labels: # 特征图内目标坐标 → 缩放后 → 底图坐标（内部重叠会保留） orig_x1 = (xc - w/2) * orig_w orig_y1 = (yc - h/2) * orig_h orig_x2 = (xc + w/2) * orig_w orig_y2 = (yc + h/2) * orig_h scaled_x1 = orig_x1 * scale_factor scaled_y1 = orig_y1 * scale_factor scaled_x2 = orig_x2 * scale_factor scaled_y2 = orig_y2 * scale_factor base_x1 = paste_x1 + scaled_x1 base_y1 = paste_y1 + scaled_y1 base_x2 = paste_x1 + scaled_x2 base_y2 = paste_y1 + scaled_y2 # 确保目标框完全在底图内（不考虑内部重叠） if base_x1 < 0 or base_y1 < 0 or base_x2 > base_w or base_y2 > base_h: valid = False break temp_target_boxes.append((base_x1, base_y1, base_x2, base_y2)) if not valid: continue # 6. 关键：检测当前特征图目标与已粘贴目标是否重叠（防跨图目标重叠） target_overlap = False for temp_box in temp_target_boxes: for existing_box in pasted_target_boxes: if calculate_iou(temp_box, existing_box) > MAX_OVERLAP_IOU: target_overlap = True skipped_target_overlap += 1 break if target_overlap: break if target_overlap: continue # 7. 粘贴特征图并记录信息（保留内部重叠） feature_img = cv2.imread(feature_img_path) if feature_img is None: continue scaled_feature_img = cv2.resize(feature_img, (scaled_w, scaled_h), interpolation=cv2.INTER_LINEAR) base_image[paste_y1:paste_y2, paste_x1:paste_x2] = scaled_feature_img # 记录标签、目标框、特征图整体区域 for (cls_id, xc, yc, w, h), temp_box in zip(feature_labels, temp_target_boxes): base_xc = (temp_box[0] + temp_box[2]) / 2 / base_w base_yc = (temp_box[1] + temp_box[3]) / 2 / base_h base_w_box = (temp_box[2] - temp_box[0]) / base_w base_h_box = (temp_box[3] - temp_box[1]) / base_h pasted_labels.append((cls_id, base_xc, base_yc, base_w_box, base_h_box)) pasted_target_boxes.append(temp_box) pasted_feature_regions.append(current_feature_region) pasted_feature_count += 1 print(f" 已成功粘贴 {pasted_feature_count}/{target_paste_count} 张特征图") # 输出统计（无内部重叠相关提示） print( f" 粘贴统计：成功{pasted_feature_count}张 → " f"跳小尺寸{skipped_small} → 跳区域重叠{skipped_region_overlap} → " f"跳目标重叠{skipped_target_overlap} → 总目标数{len(pasted_labels)}" ) return base_image, pasted_labels # --- 3. 主函数 ---（无修改） def main(): os.makedirs(OUTPUT_IMAGE_DIR, exist_ok=True) os.makedirs(OUTPUT_LABEL_DIR, exist_ok=True) try: feature_list = collect_feature_images( SOURCE_FEATURE_IMAGE_DIR, SOURCE_FEATURE_LABEL_DIR, TARGET_CLASSES ) except (FileNotFoundError, ValueError) as e: print(f"错误：{e}") return # 筛选指定数量的底图 base_files = [f for f in os.listdir(BASE_IMAGE_DIR) if f.lower().endswith(('.jpg', '.png', '.jpeg'))] total_base_count = len(base_files) if total_base_count == 0: print(f"错误：底图目录 {BASE_IMAGE_DIR} 无有效图片！") return if SELECT_BASE_IMAGE_COUNT < 0: print(f"错误：底图数量不能为负数（当前设置：{SELECT_BASE_IMAGE_COUNT}）！") return elif SELECT_BASE_IMAGE_COUNT == 0: selected_base_files = base_files print(f"\n找到 {total_base_count} 张底图，使用全部底图...\n") else: if SELECT_BASE_IMAGE_COUNT > total_base_count: selected_base_files = base_files print(f"\n警告：设置底图数（{SELECT_BASE_IMAGE_COUNT}）超过可用数（{total_base_count}），使用全部底图...\n") else: selected_base_files = random.sample(base_files, SELECT_BASE_IMAGE_COUNT) print(f"\n找到 {total_base_count} 张底图，随机选择 {SELECT_BASE_IMAGE_COUNT} 张...\n") # 处理底图 for base_filename in tqdm(selected_base_files, desc="处理底图"): base_name, ext = os.path.splitext(base_filename) base_path = os.path.join(BASE_IMAGE_DIR, base_filename) base_img = cv2.imread(base_path) if base_img is None: tqdm.write(f"\n警告：无法读取底图 {base_filename}，已跳过") continue for aug_idx in range(AUGMENTATION_FACTOR): print(f"\n底图 {base_filename}（增强序号 {aug_idx}）：", end="") base_copy = base_img.copy() pasted_img, labels = paste_feature_images_to_base(base_copy, feature_list) if not labels: tqdm.write(f"\n警告：未粘贴任何目标，已跳过该结果") continue # 保存图片和标签 output_img_name = f"{base_name}_feat_paste_{aug_idx}.jpg" output_img_path = os.path.join(OUTPUT_IMAGE_DIR, output_img_name) cv2.imwrite(output_img_path, pasted_img) output_label_name = f"{base_name}_feat_paste_{aug_idx}.txt" output_label_path = os.path.join(OUTPUT_LABEL_DIR, output_label_name) with open(output_label_path, 'w') as f: for (cls_id, x, y, w, h) in labels: f.write(f"{cls_id} {x:.6f} {y:.6f} {w:.6f} {h:.6f}\n") print(f" → 已保存（目标数：{len(labels)}）") # 统计结果 output_img_count = len([f for f in os.listdir(OUTPUT_IMAGE_DIR) if f.endswith(('.jpg', '.png'))]) output_label_count = len([f for f in os.listdir(OUTPUT_LABEL_DIR) if f.endswith('.txt')]) print(f"\n✅ 全部处理完成！") print(f" - 实际处理底图数量：{len(selected_base_files)} 张") print(f" - 生成图片：{output_img_count} 张") print(f" - 生成标签：{output_label_count} 个") print(f" - 输出路径：\n 图片 → {OUTPUT_IMAGE_DIR}\n 标签 → {OUTPUT_LABEL_DIR}") if __name__ == "__main__": main()