-
Notifications
You must be signed in to change notification settings - Fork 6
/
compute_my_dataset.py
58 lines (48 loc) · 1.64 KB
/
compute_my_dataset.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
from torchvision import transforms, datasets, models
import numpy as np
from PIL import Image
from my_dataset import NightDataSet
import os
# 计算自己数据集的均值方差
def compute_mean_and_std(dataset):
# 均值计算
mean_r = 0
mean_g = 0
mean_b = 0
for img_id in dataset.ids:
print(img_id)
img_path = os.path.join(dataset.img_root, 'sample' + img_id + '.jpg')
img = Image.open(img_path)
img = np.array(img)
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
for img_id in dataset.ids:
print(img_id)
img_path = os.path.join(dataset.img_root, 'sample' + img_id + '.jpg')
img = Image.open(img_path)
img = np.array(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 / 255.0, mean_g / 255.0, mean_r / 255.0)
std = (std_b / 255.0, std_g / 255.0, std_r / 255.0)
return mean, std
# 数据集预处理
night_root = "./" # light data set path
all_dataset = NightDataSet(night_root, transforms.ToTensor(), train_set='all.txt')
mean, std = compute_mean_and_std(all_dataset)
print(mean)
print(std)