forked from max-andr/provable-robustness-max-linear-regions
-
Notifications
You must be signed in to change notification settings - Fork 0
/
data.py
85 lines (66 loc) · 3.36 KB
/
data.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
import numpy
import numpy as np
import scipy.io
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
def get_batch_iterator(X, Y, batch_size, shuffle=False, n_batches='all'):
if n_batches == 'all':
n_batches = len(X) // batch_size
if shuffle:
x_idx = np.random.permutation(len(X))[:n_batches * batch_size]
else:
x_idx = np.arange(len(X))[:n_batches * batch_size]
for batch_idx in x_idx.reshape([len(x_idx) // batch_size, batch_size]):
batch_x, batch_y = X[batch_idx], Y[batch_idx]
yield batch_x, batch_y
def get_dataset(dataset):
# Originally, all pixels values are uint8 values in [0, 255]
train = scipy.io.loadmat('provable_robustness_max_linear_regions/datasets/{}/{}_int_train.mat'.format(dataset, dataset))
test = scipy.io.loadmat('provable_robustness_max_linear_regions/datasets/{}/{}_int_test.mat'.format(dataset, dataset))
x_train, y_train, x_test, y_test = train['images'], train['labels'], test['images'], test['labels']
y_train, y_test = dense_to_one_hot(y_train.flatten()), dense_to_one_hot(y_test.flatten())
x_train, x_test = x_train / 255.0, x_test / 255.0
if dataset == 'gts': # we drop the last partial batch
x_test, y_test = x_test[:12600], y_test[:12600]
return x_train, x_test, y_train, y_test
def dense_to_one_hot(labels_dense):
"""Convert class labels from scalars to one-hot vectors."""
num_classes = len(np.unique(labels_dense))
num_labels = labels_dense.shape[0]
index_offset = numpy.arange(num_labels) * num_classes
labels_one_hot = numpy.zeros((num_labels, num_classes))
labels_one_hot.flat[index_offset + labels_dense.ravel()] = 1
return labels_one_hot
def zero_pad(x_img, n_pad):
return np.pad(x_img, [(0, 0), (n_pad, n_pad), (n_pad, n_pad), (0, 0)], 'constant')
def augment_train(img_tensor, hps):
def augment_each(img):
if hps.random_crop:
img = tf.random_crop(img, [hps.height, hps.width, hps.n_col])
else:
img = tf.image.central_crop(img, hps.height / hps.height_pad)
if hps.dataset not in ['mnist', 'gts', 'svhn'] and hps.fl_mirroring:
img = tf.image.random_flip_left_right(img)
img = tf.image.random_brightness(img, max_delta=0.1)
img = tf.minimum(tf.maximum(img, 0.0), 1.0)
img = tf.image.random_contrast(img, lower=0.6, upper=1.4)
img = tf.minimum(tf.maximum(img, 0.0), 1.0)
return img
with tf.device('/cpu:0'):
if hps.fl_rotations:
rand_angles = tf.random_uniform([hps.batch_size], minval=-hps.max_rotate_angle, maxval=hps.max_rotate_angle)
img_tensor = tf.contrib.image.rotate(img_tensor, rand_angles)
img_tensor = tf.map_fn(augment_each, img_tensor)
if hps.gauss_noise_flag:
expected_noise_norm = 2.0 if hps.dataset in ['mnist', 'fmnist'] else 1.0
gauss_noise = tf.random_normal(tf.shape(img_tensor), stddev=expected_noise_norm / hps.n_in)
img_tensor += gauss_noise
img_tensor = tf.minimum(tf.maximum(img_tensor, 0.0), 1.0)
return img_tensor
def augment_test(img_tensor, hps):
def prepare_each(img):
img = tf.image.central_crop(img, hps.height / hps.height_pad) # crop factor is 28/32 or 32/36
return img
with tf.device('/cpu:0'):
img_tensor = tf.map_fn(prepare_each, img_tensor)
return img_tensor