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)

    # 读取类别下所有图片路径
    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("✅ 数据增强完成！")