/
utils.py
113 lines (98 loc) · 3.13 KB
/
utils.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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import torch
import numpy as np
import torchvision.transforms as trans
import math
from scipy.fftpack import dct, idct
from torch.distributions import Beta
# mean and std for different datasets
IMAGENET_SIZE = 224
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
IMAGENET_TRANSFORM = trans.Compose([
trans.Scale(256),
trans.CenterCrop(224),
trans.ToTensor()])
INCEPTION_SIZE = 299
INCEPTION_TRANSFORM = trans.Compose([
trans.Scale(342),
trans.CenterCrop(299),
trans.ToTensor()])
CIFAR_SIZE = 32
CIFAR_MEAN = [x / 255 for x in [125.3, 123.0, 113.9]]# [0.4914, 0.4822, 0.4465]
CIFAR_STD = [x / 255 for x in [63.0, 62.1, 66.7]]#[0.2023, 0.1994, 0.2010]
CIFAR_TRANSFORM = trans.Compose([
trans.ToTensor()])
MNIST_SIZE = 28
MNIST_MEAN = [0.5]
MNIST_STD = [1.0]
MNIST_TRANSFORM = trans.Compose([
trans.ToTensor()])
# reverses the normalization transformation
def invert_normalization(imgs, dataset):
if dataset == 'imagenet':
mean = IMAGENET_MEAN
std = IMAGENET_STD
elif dataset == 'cifar':
mean = CIFAR_MEAN
std = CIFAR_STD
elif dataset == 'mnist':
mean = MNIST_MEAN
std = MNIST_STD
imgs_trans = imgs.clone()
if len(imgs.size()) == 3:
for i in range(imgs.size(0)):
imgs_trans[i, :, :] = imgs_trans[i, :, :] * std[i] + mean[i]
else:
for i in range(imgs.size(1)):
imgs_trans[:, i, :, :] = imgs_trans[:, i, :, :] * std[i] + mean[i]
return imgs_trans
# applies the normalization transformations
def apply_unnormalization(imgs, dataset):
if dataset == 'imagenet':
mean = IMAGENET_MEAN
std = IMAGENET_STD
elif dataset == 'cifar':
mean = CIFAR_MEAN
std = CIFAR_STD
elif dataset == 'mnist':
mean = MNIST_MEAN
std = MNIST_STD
else:
mean = [0, 0, 0]
std = [1, 1, 1]
imgs_tensor = imgs.clone()
if dataset == 'mnist':
imgs_tensor = (imgs_tensor - mean[0]) / std[0]
else:
if imgs.dim() == 3:
for i in range(imgs_tensor.size(0)):
imgs_tensor[i, :, :] = imgs_tensor[i, :, :] * std[i]+ mean[i]
else:
for i in range(imgs_tensor.size(1)):
imgs_tensor[:, i, :, :] = imgs_tensor[:, i, :, :]* std[i]+ mean[i]
return imgs_tensor
# applies the normalization transformations
def apply_normalization(imgs, dataset):
if dataset == 'imagenet':
mean = IMAGENET_MEAN
std = IMAGENET_STD
elif dataset == 'cifar':
mean = CIFAR_MEAN
std = CIFAR_STD
elif dataset == 'mnist':
mean = MNIST_MEAN
std = MNIST_STD
else:
mean = [0, 0, 0]
std = [1, 1, 1]
imgs_tensor = imgs.clone()
if dataset == 'mnist':
imgs_tensor = (imgs_tensor - mean[0]) / std[0]
else:
if imgs.dim() == 3:
for i in range(imgs_tensor.size(0)):
imgs_tensor[i, :, :] = (imgs_tensor[i, :, :] - mean[i]) / std[i]
else:
for i in range(imgs_tensor.size(1)):
imgs_tensor[:, i, :, :] = (imgs_tensor[:, i, :, :] - mean[i]) / std[i]
return imgs_tensor