import os import cv2 import shutil import xml.etree.ElementTree as ET import random from tqdm import tqdm # 定义图片和标签文件的路径 img_path = r"D:/yolo/ultralytics-20250308/ultralytics-main/数据集处理/JPEGImages" xml_path = r"D:/yolo/ultralytics-20250308/ultralytics-main/数据集处理/Annotations" output_img_path = r"D:/yolo/ultralytics-20250308/ultralytics-main/数据集处理/outputimage" output_xml_path = r"D:/yolo/ultralytics-20250308/ultralytics-main/数据集处理/outputlabels" # 检查路径是否存在 assert os.path.exists(img_path), f"图片路径不存在: {img_path}" assert os.path.exists(xml_path), f"标签路径不存在: {xml_path}" # 创建保存增强后图片和xml文件的文件夹 os.makedirs(output_img_path, exist_ok=True) os.makedirs(output_xml_path, exist_ok=True) # 获取图片和标签文件的文件名列表 img_files = os.listdir(img_path) xml_files = os.listdir(xml_path) # 确保图片和标签文件数量匹配 assert len(img_files) == len(xml_files), "图片和标签文件数量不匹配" # 按文件名排序，确保一一对应 img_files.sort() xml_files.sort() # 定义翻转方式及其概率分布（水平、垂直、对角各占 1/3） flip_types = [0, 1, -1] # 0: 垂直翻转, 1: 水平翻转, -1: 对角翻转 # 遍历每个图片和对应的标签文件 for img_name, xml_name in tqdm(zip(img_files, xml_files), desc='Enhancing Images'): # 检查文件名是否匹配 assert img_name.split('.')[0] == xml_name.split('.')[0], f"文件名不匹配: {img_name} 和 {xml_name}" # 检查图片文件是否存在 img_file_path = os.path.join(img_path, img_name) if not os.path.exists(img_file_path): print(f"图片文件不存在: {img_name}") continue # 跳过当前图片 # 读取图片 img = cv2.imread(img_file_path) if img is None: print(f"无法加载图片: {img_name}") continue # 跳过当前图片 # 随机选择翻转方式（水平、垂直、对角各占 1/3） flip_type = random.choice(flip_types) flipped_img = cv2.flip(img, flip_type) # 检查标签文件是否存在 xml_file_path = os.path.join(xml_path, xml_name) if not os.path.exists(xml_file_path): print(f"标签文件不存在: {xml_name}") continue # 跳过当前标签 # 解析标签文件 try: tree = ET.parse(xml_file_path) root = tree.getroot() except ET.ParseError: print(f"无法解析标签文件: {xml_name}") continue # 跳过当前标签 # 检查标签文件中是否包含