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


def cal_mean_std(path: str):
    channels_sum, channels_squared_sum, nums = 0, 0, 0
    path_list = os.listdir(path)
    for img_path in path_list:
        image_path = os.path.join(path, img_path)
        image = Image.open(image_path)
        if image.mode != 'RGB':
            image = image.convert('RGB')   # convert image to RGB if it is not
        image = torch.from_numpy(np.array(image)).permute([2, 0, 1]).float()
        channels_sum += torch.mean(image, dim=[1, 2])
        channels_squared_sum += torch.mean(image**2, dim=[1, 2])
        nums += 1
    mean = channels_sum / nums
    std = (channels_squared_sum / nums - mean**2)**0.5
    return mean, std


if __name__ == '__main__':
    path = os.path.abspath(r"C:\Users\GW\Desktop\xu\data_coco\train2017")  # 数据集路径
    mean, std = cal_mean_std(path=path)
    print(f'mean : {mean}, std : {std}')