计算图片数据集的整体均值、方差和直方图

在深度学习中经常需要分析图片数据集的均值和方差以及直方图，如下代码实现了该功能，可以计算整体图片数据集的每个通道的所有像素的整体均值和方差，计算时又分包含0像素和不包含0像素，所以输出是6个均值和6个方差，之所以这样划分，原因是在深度学习中需要将所有图片缩放至同等的宽高，这时一般选用等比例拉伸，那必然带来大量的padding像素，而这些padding像素，我们一般赋予0值，所以0像素值是最大量的，远远高于第二。该代码还能输出每个通道的直方图，当然只能输出3个直方图，对应RGB，每个通道的直方图是整个数据集累加后的结果。

# 获得数据集的均值和方差，先求出每幅图的每个通道的均值和方差，然后所有图求均值，包含去零后的均值和方差
import matplotlib.pyplot as plt
import tensorflow as tf
import os
import random
import cv2
import numpy as np
import pickle


# 参数设置
IMG_SIZE = (128, 128)
dir = "./train"

# 读取图片并生成数据集类，如果不用tf的方式读取数据集，可以自己定义任何读取方式，只需要替换第25和26行的image_dataset即可
image_dataset = ""    # tf.data.dataset，是tf2.x推荐的数据读取方式，又分集成式和自定义式
# !!!所以要应用这里的代码只需要按照自己的数据集读取方式修改第25和26行代码即可，这两行本质上就是一张张把数据集的图片送过来


# 统计每个通道的直方图分布，然后计算每幅图的每个通道的均值和方差，并保存在列表中，这里分两种情况计算，分别为包含和不包含0像素值，但是直方图是累加的
mean_list, std_list = [[], []], [[], []]
hist = np.zeros((256, 3))
value = np.arange(0, 256)
num = len(file_paths_list)
for i, (images, _) in enumerate(image_dataset):  # image_dataset对应tensorflow2.x的tf.data.dataset数据读取方式，
    for j in range(8):                           # batchsize=8 这里可以更换为任何读取方式，本质上就是一张一张读取图片
        # 统计直方图
        hist_r = cv2.calcHist([images.numpy()[0]], [0], None, [256], [0, 256]).reshape((256,))
        hist_g = cv2.calcHist([images.numpy()[0]], [1], None, [256], [0, 256]).reshape((256,))
        hist_b = cv2.calcHist([images.numpy()[0]], [2], None, [256], [0, 256]).reshape((256,))
        # 计算均值和方差的共有求和数
        sum_r_h_0, sum_g_h_0, sum_b_h_0 = IMG_SIZE[0]*IMG_SIZE[1], IMG_SIZE[0]*IMG_SIZE[1], IMG_SIZE[0]*IMG_SIZE[1]
        sum_r_h,   sum_g_h,   sum_b_h   = sum_r_h_0 - hist_r[0], sum_g_h_0 - hist_g[0], sum_b_h_0 - hist_b[0]
        # 计算均值的单独求和数
        sum_r_vh_0, sum_g_vh_0, sum_b_vh_0 = np.sum(value*hist_r), np.sum(value*hist_g), np.sum(value*hist_b)
        sum_r_vh,   sum_g_vh,   sum_b_vh   = sum_r_vh_0, sum_g_vh_0, sum_b_vh_0
        # 计算均值，前后分别为包含0的统计和不包含0的统计
        mean_r_0, mean_r = sum_r_vh_0/sum_r_h_0, sum_r_vh/sum_r_h
        mean_g_0, mean_g = sum_g_vh_0/sum_g_h_0, sum_g_vh/sum_g_h
        mean_b_0, mean_b = sum_b_vh_0/sum_b_h_0, sum_b_vh/sum_b_h
        # 计算方差的单独求和数
        tmp_r, tmp_g, tmp_b = (value-mean_r_0)**2, (value-mean_g_0)**2, (value-mean_b_0)**2
        sum_r_th_0, sum_g_th_0, sum_b_th_0 = np.sum(tmp_r*hist_r), np.sum(tmp_g*hist_g), np.sum(tmp_b*hist_b)
        sum_r_th,   sum_g_th,   sum_b_th   = sum_r_th_0-tmp_r[0]*hist_r[0], sum_g_th_0-tmp_g[0]*hist_g[0], sum_b_th_0-tmp_b[0]*hist_b[0]
        # 计算方差，前后分别为包含0的统计和不包含0的统计
        std_r_0, std_r = np.sqrt(sum_r_th_0/sum_r_h_0), np.sqrt(sum_r_th/sum_r_h)
        std_g_0, std_g = np.sqrt(sum_g_th_0/sum_g_h_0), np.sqrt(sum_g_th/sum_g_h)
        std_b_0, std_b = np.sqrt(sum_b_th_0/sum_b_h_0), np.sqrt(sum_b_th/sum_b_h)
        # 记录本图的均值和方差结果
        mean_list[0].append([mean_r_0, mean_g_0, mean_b_0]) 
        mean_list[1].append([mean_r, mean_g, mean_b])
        std_list[0].append([std_r_0, std_g_0, std_b_0])
        std_list[1].append([std_r, std_g, std_b])
        # 求和直方图
        hist[:, 0] += hist_r.reshape((256,))
        hist[:, 1] += hist_g.reshape((256,))
        hist[:, 2] += hist_b.reshape((256,))
        # 打印进程
        if (i*8+j+1)%10000 == 0:
            print("finish %d -- %0.2f%%" % (i*8+j+1, (i*8+j+1)/num))

# 转换结果
mean_array = np.array(mean_list)
std_array = np.array(std_list)

# 计算总的均值和方差
mean_0, mean = mean_array[0].mean(axis=0, keepdims=True), mean_array[1].mean(axis=0, keepdims=True)
std_0, std   = std_array[0].mean(axis=0, keepdims=True), std_array[1].mean(axis=0, keepdims=True)

# 保存结果
result = {}
result["mean"] = [mean_0[0], mean[0]]
result["std"]  = [std_0[0], std[0]]
result["hist"] = hist
pickle.dump(result, open('result.pkl', 'wb'))

# 打开结果
with open('result.pkl', 'rb') as file:
    result2 = pickle.load(file)

# 打印均值方差，分别打印6个值，对应包含0和不包含0时的rgb通道
print(result2["mean"])
print(result2["std"])

# 绘制直方图，不包含0值，对应rgb通道
color = ("red", "green", "blue")
for i, color in enumerate(color):
    plt.plot(result2["hist"][1:, i], color=color)
    plt.xlim([0, 256])
plt.show()
plt.savefig('./hist.png', format='png')