# python 3.10.3
# albumentations==2.0.8
import albumentations as A
import cv2


# 人脸专属增强流水线（适配 Albumentations 2.0.8）
def a():
    return A.Compose([
        # 轻微缩放
        A.RandomScale(scale_limit=(0.1, 0.2), interpolation=cv2.INTER_LINEAR, p=0.6),
        # 缩小10% 或放大20%   缩放时的插值方式 p为概率60%执行

        # 轻微旋转
        A.Rotate(
            limit=20,  # 旋转角度范围
            interpolation=cv2.INTER_LINEAR,  # 旋转时的插值方式。 确保旋转后人脸边缘和细节不模糊（尤其适用于小角度旋转）
            border_mode=cv2.BORDER_REPLICATE,  # 旋转后边界填充方式 用图像边缘的像素值填充旋转产生的空白区域（避免黑边），保持人脸周围背景自然。
            p=0.6  # 60% 概率触发旋转，40% 概率不旋转，平衡原始姿态和旋转姿态。
        ),

        # 亮度/对比度调整
        A.RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15, p=0.6),  # 亮度调整范围
        # brightness_limit 表示图像亮度会随机增加或减少 15%（1±0.15），模拟不同光照条件（如强光、弱光）
        # contrast_limit 对比度调整范围。 表示图像对比度会随机增加或减少 15%，模拟拍摄时的光线对比度变化（如逆光导致的低对比度）。

        # 高斯噪声（修正参数）
        A.GaussNoise(
            std_range=(0.0124, 0.0215),  # 归一化后的标准差范围（对应原方差10~30）
            mean_range=(0.0, 0.0),  # 均值为0
            p=0.5
        ),
        # std_range=(0.0124, 0.0215)：噪声标准差范围（归一化后）。
        # 对应原始像素值的噪声方差约 10~30（通过 std = √var / 255 归一化），模拟相机传感器噪声或低光拍摄时的颗粒感（轻微噪声，不破坏人脸特征）。
        # mean_range=(0.0, 0.0)：噪声均值范围。  固定为 0（零均值高斯噪声），避免噪声引入偏色（如整体偏亮 / 偏暗）。

        # 高斯模糊
        A.GaussianBlur(blur_limit=(3, 5), sigma_limit=1, p=0.5),
        # blur_limit=(3, 5)：模糊核大小范围。  表示使用 3×3 或 5×5 的高斯核进行模糊，模拟相机轻微失焦或空气浑浊导致的模糊（轻微模糊，保留人脸关键特征）。

        # 色调/饱和度调整
        A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=15, p=0.5),
        # hue_shift_limit=5：色调偏移范围。 表示图像色调（颜色倾向）会随机偏移 -5~+5 度（范围 0~179），模拟不同光源色温（如暖光、冷光）导致的颜色变化。
        # sat_shift_limit=15：饱和度偏移范围。 表示颜色饱和度会随机增加或减少 15%，模拟不同光线条件下的色彩鲜艳度变化（如阴天饱和度低，晴天饱和度高）。
			

        A.HorizontalFlip(p=0.3),  # 30% 概率水平翻转
        A.PadIfNeeded(p=0.2), # 填充到固定尺寸，模拟人脸靠近边缘的情况
        A.RandomGamma(p=0.3),  # 伽马值在 0.8~1.2 之间调整
        A.Sharpen(alpha=(0.1, 0.3), lightness=(0.8, 1.2), p=0.2),  # 轻微锐化
        A.CoarseDropout(p=0.1),  # 随机遮挡少量像素，模拟局部噪点或遮挡
        A.ElasticTransform(p=0.1),  # 轻微弹性变形，保留人脸整体结构

        # 固定输出尺寸（修正参数）
        A.Resize(height=640, width=640),  # 关键修正
        # 强制所有增强后的图像统一为 640×640 像素，确保后续人脸检测（如 InsightFace）的输入尺寸一致，避免因尺寸变化导致检测精度下降。

    ])


# 用法：对单张cv2图像增强
def augment_face(img):
    aug = a()
    # Albumentations默认处理RGB，若输入是BGR需转换
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    augmented = aug(image=img_rgb)
    img_aug_rgb = augmented["image"]
    # 转回BGR（若后续用cv2处理）
    return cv2.cvtColor(img_aug_rgb, cv2.COLOR_RGB2BGR)


# 生成3个增强变体（原图+2增强图）
img = cv2.imread(r"C:\Users\123\Desktop\face\bo.png")
img = augment_face(img)
cv2.imshow('', img)
cv2.waitKey(0)