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}')