51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
import os
|
|
import random
|
|
from PIL import Image
|
|
from torchvision import transforms
|
|
|
|
# 原始数据目录
|
|
input_root = "../dataset/train"
|
|
|
|
# 增强后保存目录
|
|
output_root = "../dataset/train_aug"
|
|
|
|
# 每类目标张数
|
|
target_num = 500
|
|
|
|
# 定义数据增强
|
|
transform = transforms.Compose([
|
|
transforms.RandomHorizontalFlip(p=0.5),
|
|
transforms.RandomRotation(15),
|
|
transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.02),
|
|
transforms.RandomResizedCrop(size=(224, 224), scale=(0.8, 1.0)),
|
|
])
|
|
|
|
os.makedirs(output_root, exist_ok=True)
|
|
|
|
# 遍历每个类别文件夹
|
|
for class_name in os.listdir(input_root):
|
|
input_dir = os.path.join(input_root, class_name)
|
|
output_dir = os.path.join(output_root, class_name)
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
if class_name == "手撕包菜":
|
|
# 读取类别下所有图片路径
|
|
img_files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.jpg', '.png', '.jpeg'))]
|
|
img_paths = [os.path.join(input_dir, f) for f in img_files]
|
|
|
|
print(f"类别 {class_name} 原始图片数: {len(img_paths)}")
|
|
|
|
count = 0
|
|
while count < target_num:
|
|
img_path = random.choice(img_paths)
|
|
img = Image.open(img_path).convert("RGB")
|
|
# 生成增强图
|
|
aug_img = transform(img)
|
|
# 保存
|
|
save_path = os.path.join(output_dir, f"aug_{count:03d}.jpg")
|
|
aug_img.save(save_path)
|
|
count += 1
|
|
|
|
print(f"类别 {class_name} 已扩充到 {target_num} 张，保存于 {output_dir}")
|
|
|
|
print("✅ 数据增强完成！")
|