从文件夹中随机选择一定数量的图像,然后对每个选定的图像进行一次随机的数据增强变换。
import os
import random
import cv2
import numpy as np
from PIL import Image, ImageEnhance, ImageOps
# 定义各种数据增强方法
def random_rotate(image, angle_range=(-30, 30)):
angle = random.uniform(angle_range[0], angle_range[1])
(h, w) = image.shape[:2]
center = (w // 2, h // 2)
M = cv2.getRotationMatrix2D(center, angle, 1.0)
rotated = cv2.warpAffine(image, M, (w, h), borderMode=cv2.BORDER_REFLECT)
return rotated
def random_translate(image, translate_range=(-50, 50)):
tx = random.randint(translate_range[0], translate_range[1])
ty = random.randint(translate_range[0], translate_range[1])
(h, w) = image.shape[:2]
M = np.float32([[1, 0, tx], [0, 1, ty]])
translated = cv2.warpAffine(image, M, (w, h), borderMode=cv2.BORDER_REFLECT)
return translated
def random_flip(image):
flip_code = random.choice([-1, 0, 1])
flipped = cv2.flip(image, flip_code)
return flipped
def random_scale(image, scale_range=(0.8, 1.2)):
scale = random.uniform(scale_range[0], scale_range[1])
(h, w) = image.shape[:2]
new_dim = (int(w * scale), int(h * scale))
scaled = cv2.resize(image, new_dim, interpolation=cv2.INTER_LINEAR)
return scaled
def random_crop(image, crop_size=(224, 224)):
(h, w) = image.shape[:2]
if crop_size[0] > h or crop_size[1] > w:
# 当裁剪尺寸大于图像尺寸时,抛出异常或调整裁剪尺寸
raise ValueError("Crop size is larger than image size.")
top = random.randint(0, h - crop_size[0])
left = random.randint(0, w - crop_size[1])
cropped = image[top:top+crop_size[0], left:left+crop_size[1]]
return cropped
def random_color_jitter(image):
pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
color_jitter = ImageEnhance.Color(pil_image).enhance(random.uniform(0.6, 1.4))
contrast_jitter = ImageEnhance.Contrast(color_jitter).enhance(random.uniform(0.5, 1.5))
brightness_jitter = ImageEnhance.Brightness(contrast_jitter).enhance(random.uniform(0.6, 1.4))
sharpness_jitter = ImageEnhance.Sharpness(brightness_jitter).enhance(random.uniform(0.6, 1.4))
jittered = cv2.cvtColor(np.array(sharpness_jitter), cv2.COLOR_RGB2BGR)
return jittered
def random_add_noise(image):
row, col, ch = image.shape
mean = 0
var = 0.1
sigma = var ** 0.5
gauss = np.random.normal(mean, sigma, (row, col, ch))
gauss = gauss.reshape(row, col, ch)
noisy = image + gauss
return np.clip(noisy, 0, 255).astype(np.uint8)
# 数据增强主函数
def augment_random_images(src_folder, dst_folder, num_images_to_select, num_augmentations_per_image):
if not os.path.exists(dst_folder):
os.makedirs(dst_folder)
# 获取所有图像文件名
all_filenames = [f for f in os.listdir(src_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
# 如果选择的图像数量大于总图像数量,则只处理全部图像
num_images_to_process = min(num_images_to_select, len(all_filenames))
# 随机选择图像
selected_filenames = random.sample(all_filenames, num_images_to_process)
# 创建一个增强方法列表
augmentation_methods = [
random_rotate,
#random_translate,
random_flip,
random_scale,
#random_crop,
random_color_jitter,
random_add_noise
]
for filename in selected_filenames:
img_path = os.path.join(src_folder, filename)
image = cv2.imread(img_path)
for i in range(num_augmentations_per_image):
# 随机选择一种增强方法
augmentation_method = random.choice(augmentation_methods)
# 应用选中的增强方法
augmented_img = augmentation_method(image)
# 保存增强后的图像
base_name, ext = os.path.splitext(filename)
save_path = os.path.join(dst_folder, f"{base_name}_aug_{i}{ext}")
cv2.imwrite(save_path, augmented_img)
if __name__ == "__main__":
src_folder = 'path/to/source/folder' # 替换为你的源文件夹路径
dst_folder = 'path/to/destination/folder' # 替换为你要保存增强图像的文件夹路径
num_images_to_select = 10 # 从源文件夹中随机选择的图像数量
num_augmentations_per_image = 5 # 每张图像生成的增强图像数量
augment_random_images(src_folder, dst_folder, num_images_to_select, num_augmentations_per_image)
print(f"图像增强完成,增强后的图像已保存到 {dst_folder}")
说明
- 随机选择图像:从源文件夹中随机选择num_images_to_select数量的图像。
- 随机选择一种增强方法:对于每张选定的图像,随机选择一种数据增强方法。
- 应用增强方法:对每张选定的图像应用所选的增强方法。
- 保存增强后的图像:将增强后的图像保存到目标文件夹中。
参数
•src_folder:源文件夹路径。
•dst_folder:目标文件夹路径。
•num_images_to_select:从源文件夹中随机选择的图像数量。
•num_augmentations_per_image:每张选定的图像生成的增强图像数量。
请确保将src_folder和dst_folder变量设置为您实际使用的文件夹路径,并根据需要调整num_images_to_select和num_augmentations_per_image的值。运行这段代码后,将得到从源文件夹中随机选择的图像,并对这些图像进行了随机的数据增强变换。