【猫狗数据集】计算数据集的平均值和方差

import os
import cv2
import numpy as np
from torch.utils.data import Dataset
from PIL import Image
import torchvision
import time
from time import time 
from tqdm import tqdm

def compute_mean_and_std(dataset):
    # 输入PyTorch的dataset，输出均值和标准差
    mean_r = 0
    mean_g = 0
    mean_b = 0
    print("计算均值>>>")
    for img_path, _ in tqdm(dataset,ncols=80):
      img=Image.open(img_path)
      img = np.asarray(img) # change PIL Image to numpy array
      mean_b += np.mean(img[:, :, 0])
      mean_g += np.mean(img[:, :, 1])
      mean_r += np.mean(img[:, :, 2])

    mean_b /= len(dataset)
    mean_g /= len(dataset)
    mean_r /= len(dataset)

    diff_r = 0
    diff_g = 0
    diff_b = 0

    N = 0
    print("计算方差>>>")
    for img_path, _ in tqdm(dataset,ncols=80):
      img=Image.open(img_path)
      img = np.asarray(img)
      diff_b += np.sum(np.power(img[:, :, 0] - mean_b, 2))
      diff_g += np.sum(np.power(img[:, :, 1] - mean_g, 2))
      diff_r += np.sum(np.power(img[:, :, 2] - mean_r, 2))

      N += np.prod(img[:, :, 0].shape)

    std_b = np.sqrt(diff_b / N)
    std_g = np.sqrt(diff_g / N)
    std_r = np.sqrt(diff_r / N)

    mean = (mean_b.item() / 255.0, mean_g.item() / 255.0, mean_r.item() / 255.0)
    std = (std_b.item() / 255.0, std_g.item() / 255.0, std_r.item() / 255.0)
    return mean, std
path = "/content/drive/My Drive/colab notebooks/data/dogcat"
train_path=path+"/train"
test_path=path+"/test"
val_path=path+'/val'
train_data = torchvision.datasets.ImageFolder(train_path)
val_data = torchvision.datasets.ImageFolder(val_path)
test_data = torchvision.datasets.ImageFolder(test_path)
#train_mean,train_std=compute_mean_and_std(train_data.imgs)
time_start =time()
val_mean,val_std=compute_mean_and_std(val_data.imgs)
time_end=time()
print("验证集计算消耗时间：", round(time_end - time_start, 4), "s")
#test_mean,test_std=compute_mean_and_std(test_data.imgs)
#print("训练集的平均值：{}，方差：{}".format(train_mean,train_std))
print("验证集的平均值：{}".format(val_mean))
print("验证集的方差：{}".format(val_mean))
#print("测试集的平均值：{}，方差：{}".format(test_mean,test_std))

train_transform = torchvision.transforms.Compose([
    torchvision.transforms.RandomResizedCrop(size=224,
                                             scale=(0.08, 1.0)),
    torchvision.transforms.RandomHorizontalFlip(),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
                                     std=(0.229, 0.224, 0.225)),
 ])
 val_transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize(256),
    torchvision.transforms.CenterCrop(224),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
                                     std=(0.229, 0.224, 0.225)),
])

import numpy as np
import cv2
import random
 
# calculate means and std
train_txt_path = './train_val_list.txt'
 
CNum = 10000   # 挑选多少图片进行计算
 
img_h, img_w = 224, 224
imgs = np.zeros([img_w, img_h, 3, 1])
means, stdevs = [], []
 
with open(train_txt_path, 'r') as f:
  lines = f.readlines()
  random.shuffle(lines)  # shuffle , 随机挑选图片
 
  for i in tqdm_notebook(range(CNum)):
    img_path = os.path.join('./train', lines[i].rstrip().split()[0])
 
    img = cv2.imread(img_path)
    img = cv2.resize(img, (img_h, img_w))
    img = img[:, :, :, np.newaxis]
    
    imgs = np.concatenate((imgs, img), axis=3)
#     print(i)
 
imgs = imgs.astype(np.float32)/255.
 
 
for i in tqdm_notebook(range(3)):
  pixels = imgs[:,:,i,:].ravel() # 拉成一行
  means.append(np.mean(pixels))
  stdevs.append(np.std(pixels))
 
# cv2 读取的图像格式为BGR，PIL/Skimage读取到的都是RGB不用转
means.reverse() # BGR --> RGB
stdevs.reverse()
 
print("normMean = {}".format(means))
print("normStd = {}".format(stdevs))
print('transforms.Normalize(normMean = {}, normStd = {})'.format(means, stdevs))