#import PIL, os
from tqdm import tqdm
import numpy as np
import torch
from torch.utils.data import Dataset
import glob
import cv2
import os
# https://gist.github.com/spirosdim/79fc88231fffec347f1ad5d14a36b5a8
class ImgDataset(Dataset):
def __init__(self, data_dir, extension='.jpg', transform=None):
self.data_dir = data_dir
self.image_files = []
for it in glob.glob(os.path.join(data_dir, "*")):
print(it)
self.image_files += list(glob.glob(os.path.join(it, '*' + extension)))
#self.image_files = [f for f in os.listdir(data_dir) if f.endswith(extension)]
self.transform = transform
print(f"image size: {len(self.image_files)}")
def __len__(self):
return len(self.image_files)
def __getitem__(self, index):
im = cv2.imread(self.image_files[index])
im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
im = im.astype(float) / 255.
image = torch.from_numpy(im)
#image = torch.from_numpy(np.array(PIL.Image.open(self.image_files[index]),dtype=np.float32))
if self.transform:
image = self.transform(image)
return image
def compute_mean_std(dataloader):
'''
We assume that the images of the dataloader have the same height and width
source: https://github.com/aladdinpersson/Machine-Learning-Collection/blob/master/ML/Pytorch/Basics/pytorch_std_mean.py
'''
# var[X] = E[X**2] - E[X]**2
channels_sum, channels_sqrd_sum, num_batches = 0, 0, 0
for batch_images in tqdm(dataloader): # (B,H,W,C)
channels_sum += torch.mean(batch_images, dim=[0, 1, 2])
channels_sqrd_sum += torch.mean(batch_images ** 2, dim=[0, 1, 2])
num_batches += 1
mean = channels_sum / num_batches
std = (channels_sqrd_sum / num_batches - mean ** 2) ** 0.5
return mean, std
if __name__ == "__main__":
dataset = ImgDataset('data/train')
train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=False,
num_workers=16, pin_memory=True)
# output
total_mean, total_std = compute_mean_std(train_loader)
print('mean (RGB): ' + str(total_mean))
print('std (RGB): ' + str(total_std))
python计算图片数据集均值方差
于 2024-03-25 10:26:51 首次发布