[Pytorch] custom dataset의 mean, std 구하기

import numpy as np
import matplotlib.pyplot as plt
from pycocotools.coco import COCO
import torch
from torchvision import datasets
from tqdm.notebook import tqdm

def image_stats(data):
    total_images = len(data)
    total_pixels = 0
    pixel_sum = np.zeros(3)
    pixel_squared_sum = np.zeros(3)

    for i in range(total_images):
        # 이미지 데이터를 numpy 객체로 변환하고 정규화(0~1)
        img, _ = data[i]
        img_np = np.asarray(img) / 255.0
        total_pixels += img_np.size // img_np.shape[-1]

        # 픽셀값의 합과 제곱의 합을 계산
        pixel_sum += np.sum(img_np, axis=(0, 1))
        pixel_squared_sum += np.sum(np.square(img_np), axis=(0, 1))

    # 평균과 표준편차 계산 및 반환
    mean = pixel_sum / total_pixels
    std = np.sqrt(pixel_squared_sum / total_pixels - np.square(mean))
    return mean, std
    
    
train_data = datasets.CocoDetection('train데이터셋의 위치', 'instances_train.json파일의 위치')
#test_data = datasets.CocoDetection('test데이터셋의 위치', 'instances_test.json파일의 위치')

mean, std = image_stats(train_data)
print("Mean:", mean)
print("Std:", std)