【數據增強】精細化貼圖數據增強

1.任務背景

假設我有100個蘋果的照片，我需要把這些照片粘貼到傳送帶照片上，模擬“傳送帶蘋果檢測”場景。
這種貼圖的方式更加合理一些，因為yolo之類的mosaic貼圖，會把圖像弄的非常支離破碎。
現在我需要隨機選擇幾張蘋果圖像，每張蘋果圖像至少使用x次，并且保證蘋果（新蘋果之間、新舊標注信息之間）不重疊，且蘋果大小范圍可以自由指定。
效果圖如下（人工核查，xml也是正確的，我粘貼的是電力部件）。在真實使用中，如果需要確保信粘貼內容與已有標簽不重復，請在在這里插入圖片描述
）：

2.具體邏輯

2.1 項目描述

這是一個用于計算機視覺任務（如目標檢測）的智能數據增強腳本。它通過將小物體圖像（Patches）以復雜且真實的方式粘貼到大型背景圖像上，來批量生成高質量的訓練數據集。腳本的核心是“場景合成”，旨在創建包含多個、尺寸合理且互不重疊物體的復雜圖像，從而有效提升模型的魯棒性和泛化能力。

2.2 整體流程

腳本的運行流程遵循一個清晰的、基于統計學控制的“事件驅動”模型：

計算任務總量: 首先，根據用戶設定的參數（小圖總數、期望重復次數、每背景粘貼數），腳本會計算出需要生成的增強圖片總數。
循環生成: 程序會循環執行所計算出的總次數。在每一次循環中，它會獨立地完成一次“場景合成”操作。
場景合成:
- 選取素材: 隨機選擇一張背景圖和指定數量的隨機小圖。
- 智能調整與放置: 對每一個選中的小圖，依次進行：
  - 動態縮放: 根據相對于背景圖的“最大/最小面積百分比”要求，自動縮放小圖，確保尺寸合理。
  - 碰撞檢測: 為小圖尋找一個隨機的、且不與任何已放置小圖重疊的位置。
  - 數據增強: 對小圖進行隨機的水平或垂直翻轉。
  - 粘貼與記錄: 將處理后的小圖粘貼到背景圖上，并同步更新XML標注信息。
- 保存輸出: 將最終合成的圖片和包含所有新物體標注的XML文件，以唯一編號保存到輸出文件夾。

2.3主要功能

場景化組合: 能在單一背景上粘貼多個隨機物體，模擬復雜的真實世界場景。
動態尺寸調整: 摒棄了固定的像素限制，采用相對于背景的面積百分比來約束貼圖大小，使其能智能適應任意尺寸的背景圖。
防重疊粘貼: 核心亮點功能。通過碰撞檢測算法確保粘貼的物體之間（包括原有的標注信息）互不重疊，顯著提升了生成數據的質量和真實感。
可控的隨機性: 用戶可以通過參數精確控制最終生成數據集的總量，并使每個物體的平均使用次數在統計上趨于穩定。
自動化標注: 在生成圖像的同時，會自動創建和更新對應的PASCAL VOC格式的XML標注文件，省去手動標注的繁瑣工作。

2.4 參數說明

FOLDER_A, FOLDER_B, OUTPUT...: 用于定義素材和輸出結果的路徑。
OBJECT_NAME: 指定粘貼的物體在XML文件中被稱作什么類別名。
NUM_PATCHES_PER_BG: **（核心）**設定每張生成的圖片上要粘貼多少個小圖/物體。
REPEATS_PER_PATCH: **（核心）**設定數據集中每一種小圖期望被重復使用的平均次數，用于計算總生成量。
MAX_AREA_PERCENTAGE: 定義貼圖相對于背景的最大允許面積（例如 0.5 代表50%）。
MIN_AREA_PERCENTAGE: 定義貼圖相對于背景的最小允許面積（例如 0.05 代表5%），小于此值會被自動放大。
MAX_PLACEMENT_TRIES: 在為小圖尋找不重疊位置時的最大嘗試次數，這是一個防止在擁擠場景下無限循環的安全設置。

3.代碼實現

import os
import random
import xml.etree.ElementTree as ET
from PIL import Image
from tqdm import tqdm
import copy
import math
import cv2# --- 配置區 ---# 輸入文件夾
FOLDER_A = r'E:\data\baidu_pic\aaa'  # 存放小圖像的文件夾
FOLDER_B = r'E:\data\baidu_pic\background'  # 存放帶XML的大圖像的文件夾# 輸出文件夾 (如果不存在，腳本會自動創建)
OUTPUT_IMAGES_FOLDER = r'E:\data\baidu_pic\hecheng\iamges'
OUTPUT_ANNOTATIONS_FOLDER = r'E:\data\baidu_pic\hecheng\xlms'# 要在XML中添加的Object名稱
OBJECT_NAME = 'rdg'# --- 關鍵參數 ---
# 1. 為每張大圖粘貼多少個小圖
NUM_PATCHES_PER_BG = 3# 2. 數據集中每張小圖期望重復使用的次數
# 這將決定“小圖池”的大小
REPEATS_PER_PATCH = 2# 3.新增：相對面積限制 
# 使用基于背景圖面積的百分比來控制貼圖的最終尺寸
MAX_AREA_PERCENTAGE = 0.01  # 貼圖面積不得超過背景的1%
MIN_AREA_PERCENTAGE = 0.005 # 貼圖面積不得小于背景的0.5%# 4.支持的圖像文件擴展名
SUPPORTED_IMAGE_FORMATS = ['.jpg', '.jpeg', '.png', '.bmp']# 5.為防止因空間不足而無限循環，設定為每個貼圖尋找不重疊位置的最大嘗試次數
MAX_PLACEMENT_TRIES = 100# --- 輔助函數 ---def is_overlapping(box1, box2):"""檢查兩個邊界框是否重疊。 box = (xmin, ymin, xmax, ymax)"""if box1[2] < box2[0] or box2[2] < box1[0]:return Falseif box1[3] < box2[1] or box2[3] < box1[1]:return Falsereturn Truedef get_b_image_basenames(folder_path):filenames = []for f in os.listdir(folder_path):basename, ext = os.path.splitext(f)if ext.lower() in SUPPORTED_IMAGE_FORMATS and os.path.exists(os.path.join(folder_path, basename + '.xml')):filenames.append(basename)return list(set(filenames))def update_xml_annotation(xml_root, object_name, xmin, ymin, xmax, ymax):obj = ET.SubElement(xml_root, 'object')ET.SubElement(obj, 'name').text = object_nameET.SubElement(obj, 'pose').text = 'Unspecified'ET.SubElement(obj, 'truncated').text = '0'ET.SubElement(obj, 'difficult').text = '0'bndbox = ET.SubElement(obj, 'bndbox')ET.SubElement(bndbox, 'xmin').text = str(int(xmin))ET.SubElement(bndbox, 'ymin').text = str(int(ymin))ET.SubElement(bndbox, 'xmax').text = str(int(xmax))ET.SubElement(bndbox, 'ymax').text = str(int(ymax))def find_image_ext(folder, basename):for ext in SUPPORTED_IMAGE_FORMATS:if os.path.exists(os.path.join(folder, basename + ext)):return extreturn None# --- 主邏輯區 ---def main():print("--- 開始執行數據增強任務 (帶完全防重疊邏輯) ---")os.makedirs(OUTPUT_IMAGES_FOLDER, exist_ok=True)os.makedirs(OUTPUT_ANNOTATIONS_FOLDER, exist_ok=True)all_small_images = [f for f in os.listdir(FOLDER_A) if os.path.splitext(f)[1].lower() in SUPPORTED_IMAGE_FORMATS]if not all_small_images:print(f"錯誤: 文件夾 '{FOLDER_A}' 中沒有找到任何圖像文件。")returnbackground_basenames = get_b_image_basenames(FOLDER_B)if not background_basenames:print(f"錯誤: 文件夾 '{FOLDER_B}' 中沒有找到任何帶有XML配對的圖像文件。")returnnum_total_patches = len(all_small_images)if NUM_PATCHES_PER_BG <= 0:print("錯誤: NUM_PATCHES_PER_BG 必須大于0。")returntotal_operations = int((num_total_patches * REPEATS_PER_PATCH) / NUM_PATCHES_PER_BG)print(f"將生成 {total_operations} 張增強圖片。")for i in tqdm(range(total_operations), desc="生成增強圖片中"):try:bg_basename = random.choice(background_basenames)bg_ext = find_image_ext(FOLDER_B, bg_basename)if not bg_ext: continuebg_image_path = os.path.join(FOLDER_B, bg_basename + bg_ext)xml_path = os.path.join(FOLDER_B, bg_basename + '.xml')background_image_pil = Image.open(bg_image_path).convert("RGBA")tree = ET.parse(xml_path)xml_root = tree.getroot()bg_width, bg_height = background_image_pil.sizebackground_area = bg_width * bg_height# --- 核心改動：預加載原始XML中的所有物體邊界框 ---placed_boxes = []for obj in xml_root.findall('object'):try:bndbox = obj.find('bndbox')# 將XML中的坐標文本轉換為整數xmin = int(float(bndbox.find('xmin').text))ymin = int(float(bndbox.find('ymin').text))xmax = int(float(bndbox.find('xmax').text))ymax = int(float(bndbox.find('ymax').text))placed_boxes.append((xmin, ymin, xmax, ymax))except (AttributeError, ValueError) as e:print(f"\n警告: 解析背景'{bg_basename}'的XML時，有對象格式不正確，已跳過。錯誤: {e}")continue# ----------------------------------------------------patches_to_paste = random.choices(all_small_images, k=NUM_PATCHES_PER_BG)for small_img_filename in patches_to_paste:small_image_path = os.path.join(FOLDER_A, small_img_filename)patch_cv_image = cv2.imread(small_image_path, cv2.IMREAD_UNCHANGED)if patch_cv_image is None: continueh, w = patch_cv_image.shape[:2]patch_area = w * hmax_allowed_area = background_area * MAX_AREA_PERCENTAGEif patch_area > max_allowed_area:scale_factor = math.sqrt(max_allowed_area / patch_area)new_w, new_h = int(w * scale_factor), int(h * scale_factor)patch_cv_image = cv2.resize(patch_cv_image, (new_w, new_h), interpolation=cv2.INTER_AREA)h, w, patch_area = new_h, new_w, new_w * new_hmin_required_area = background_area * MIN_AREA_PERCENTAGEif patch_area < min_required_area:scale_factor = math.sqrt(min_required_area / patch_area)new_w, new_h = int(w * scale_factor), int(h * scale_factor)if new_w >= bg_width or new_h >= bg_height: continuepatch_cv_image = cv2.resize(patch_cv_image, (new_w, new_h), interpolation=cv2.INTER_CUBIC)if len(patch_cv_image.shape) == 3 and patch_cv_image.shape[2] == 4:patch_cv_image = cv2.cvtColor(patch_cv_image, cv2.COLOR_BGRA2RGBA)else:patch_cv_image = cv2.cvtColor(patch_cv_image, cv2.COLOR_BGR2RGB)patch_image = Image.fromarray(patch_cv_image)if random.choice([True, False]): patch_image = patch_image.transpose(Image.FLIP_LEFT_RIGHT)if random.choice([True, False]): patch_image = patch_image.transpose(Image.FLIP_TOP_BOTTOM)patch_width, patch_height = patch_image.sizefor _ in range(MAX_PLACEMENT_TRIES):paste_x = random.randint(0, bg_width - patch_width)paste_y = random.randint(0, bg_height - patch_height)candidate_box = (paste_x, paste_y, paste_x + patch_width, paste_y + patch_height)is_valid_placement = Truefor existing_box in placed_boxes:if is_overlapping(candidate_box, existing_box):is_valid_placement = Falsebreakif is_valid_placement:background_image_pil.paste(patch_image, (paste_x, paste_y), patch_image)update_xml_annotation(xml_root, OBJECT_NAME, paste_x, paste_y, candidate_box[2], candidate_box[3])placed_boxes.append(candidate_box)breakelse:print(f"\n警告: 在嘗試 {MAX_PLACEMENT_TRIES} 次后，未能為小圖 '{small_img_filename}' 找到不重疊的位置。跳過此小圖。")new_basename = f"augmented_output_{i+1}"output_image_path = os.path.join(OUTPUT_IMAGES_FOLDER, new_basename + bg_ext)output_xml_path = os.path.join(OUTPUT_ANNOTATIONS_FOLDER, new_basename + ".xml")background_image_pil.convert("RGB").save(output_image_path)tree.write(output_xml_path, encoding='utf-8')except Exception as e:print(f"\n在生成第 {i+1} 張圖片時發生未知錯誤: {e}")continueprint(f"\n--- 所有任務已完成！---")if __name__ == '__main__':main()