pytorch 跑深度学习数据集的时候,需要先将数据归一化,可以让网络更好的收敛。一般的均值和方差都是利用(灰度值 - 0.5 )/ 0.5
本章提供一个计算均值和方差的代码,供以后使用
代码如下:
import numpy as np
import os
from PIL import Image
import cv2
# 输入 PyTorch的 dataset, 输出均值和标准差
def compute_mean_and_std(path):
img_path = [os.path.join(path,i) for i in os.listdir(path)] # 所有数据的绝对路径
mean_r,mean_g,mean_b = 0,0,0
for i in img_path:
img = np.array(Image.open(i).convert('RGB'))
img = cv2.cvtColor(img,cv2.COLOR_RGB2BGR)
mean_b += np.mean(img[:, :, 0]) # 所有数据的灰度值相加
mean_g += np.mean(img[:, :, 1])
mean_r += np.mean(img[:, :, 2])
mean_b /= len(path) # 平均灰度
mean_g /= len(path)
mean_r /= len(path)
N = 0
diff_r, diff_g, diff_b = 0, 0, 0
for i in img_path:
img = np.array(Image.open(i).convert('RGB'))
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
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 np.round(mean,3),np.round(std,3)
if __name__ == '__main__':
root = './data' # 数据的根目录
mean,std = compute_mean_and_std(path=root)
print(mean)
print(std)
均值和方差计算就是简单利用数学方法计算
这里需要注意的是,有些数据集是单通道的,这样代码计算3通道就会报错。所以这里用PIL和opencv 将图像转换
传入数据的时候,只需要传入父目录即可,demo脚本和data在同一路径下