import cv2
import numpy as np
import os
import glob# 数据增强函数
def augment_data(img):rows,cols,_ = img.shape# 水平翻转图像if np.random.random() > 0.5:img = cv2.flip(img, 1)img_name = os.path.splitext(save_path)[0] + "_flip.png"cv2.imwrite(img_name, img)print("Saved augmented image:", img_name)# 随机缩放图像scale = np.random.uniform(0.9, 1.1)M = cv2.getRotationMatrix2D((cols/2, rows/2), 0, scale)img_transformed = cv2.warpAffine(img, M, (cols, rows))img_name = os.path.splitext(save_path)[0] + "_transform.png"cv2.imwrite(img_name, img_transformed)print("Saved augmented image:", img_name)# 随机旋转图像angle = np.random.randint(-10, 10)M = cv2.getRotationMatrix2D((cols/2, rows/2), angle, 1)img_rotated = cv2.warpAffine(img, M, (cols, rows))img_name = os.path.splitext(save_path)[0] + "_rotated.png"cv2.imwrite(img_name, img_rotated)print("Saved augmented image:", img_name)# 添加高斯噪音mean = 0std = np.random.uniform(5, 15)noise = np.zeros(img.shape, np.float32)cv2.randn(noise, mean, std)noise = np.uint8(noise)img_noisy = cv2.add(img, noise)img_name = os.path.splitext(save_path)[0] + "_noisy.png"cv2.imwrite(img_name, img_noisy)print("Saved augmented image:", img_name)# 随机调整对比度和亮度alpha = np.random.uniform(0.8, 1.2)beta = np.random.randint(-10, 10)img_contrast = cv2.convertScaleAbs(img, alpha=alpha, beta=beta)img_name = os.path.splitext(save_path)[0] + "_contrast.png"cv2.imwrite(img_name, img_contrast)print("Saved augmented image:", img_name)return img# 读取 data 文件夹中的所有图片，并进行数据增强
data_dir = "data"
save_dir = "result"
if not os.path.exists(save_dir):os.makedirs(save_dir)# 使用 glob 库来遍历 data 文件夹中所有图像
for img_path in glob.glob(os.path.join(data_dir, "*.png")):img = cv2.imread(img_path)# 获取保存增强后的图片文件名img_name = os.path.basename(img_path)save_path = os.path.join(save_dir, img_name)# 数据增强augment_data(img)# 保存原始图片cv2.imwrite(save_path, img)print("Saved original image:", save_path)

XML格式数据集转TXT（YOLO）_xml转txt-CSDN博客