from PIL import Image
import os
import numpy as np
from tqdm import tqdm

def calculate_pic_std_and_mean(dataset_dir = r'./My_Data/A_Ori'):
    # 获取所有jpg图像文件
    image_files = [os.path.join(dataset_dir, filename) for filename in os.listdir(dataset_dir) if filename.lower().endswith(('.jpg', '.png', '.tiff', '.jpeg', '.bmp'))]

    # 初始化用于存储累积的像素值
    sum_pixels_normalized = np.zeros(3)
    sum_squared_pixels_normalized = np.zeros(3)
    num_pixels = 0

    # 使用tqdm创建一个进度条
    for image_file in tqdm(image_files, desc="Calculating mean and std"):
        image = Image.open(image_file).convert('RGB')  # 确保图像为RGB
        image = np.array(image)  # 原图像像素范围[0, 255]
        
        # 归一化到[0, 1]范围
        image_normalized = image / 255.0

        # 累积归一化像素值和归一化像素平方值
        sum_pixels_normalized += np.sum(image_normalized, axis=(0, 1))  # 按通道累积
        sum_squared_pixels_normalized += np.sum(image_normalized ** 2, axis=(0, 1))  # 按通道累积像素平方值
        num_pixels += image.shape[0] * image.shape[1]  # 累积总像素数

    # 计算整个数据集的归一化后的均值
    mean_normalized = sum_pixels_normalized / num_pixels

    # 计算整个数据集的归一化后的标准差
    variance_normalized = sum_squared_pixels_normalized / num_pixels - mean_normalized ** 2
    variance_normalized = np.maximum(variance_normalized, 0)  # 防止负数
    std_normalized = np.sqrt(variance_normalized)

    # 反归一化回[0, 255]范围
    mean = mean_normalized * 255
    std = std_normalized * 255

    print(f"\033[93m计算得图片均值-Mean: {mean}  计算得图片方差-Std: {std}\033[0m")
    
    return mean, std

if __name__ == '__main__':
    # 示例调用
    calculate_pic_std_and_mean()
