-
Notifications
You must be signed in to change notification settings - Fork 9
/
data_transform.py
88 lines (71 loc) · 3.16 KB
/
data_transform.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
import random
import numpy as np
import matplotlib.pyplot as plt
from skimage.transform import AffineTransform
from skimage.transform import warp
from transforming_autoencoders.utils.data_structures import TransformingAutoencoderExample
def get_random_affine_matrix(sigma, max_translation):
"""
Get a random affine transformation matrix
Parameters
----------
sigma: float
Parametrizes the gaussian noise added to transformation matrix
max_translation:
Maximum translation allowed in the transformation matrix
Returns
-------
matrix: ndarray
3x3 matrix representing a random affine transformation
"""
R = np.eye(2) + sigma * np.random.normal(0, 1, size=[2, 2]) # rotation and scale
T = np.random.uniform(-max_translation, max_translation, size=[2, 1]) # translation
H = np.array([[0., 0., 1.]]) # homogeneous part
return np.concatenate([np.concatenate([R, T], axis=1), H], axis=0)
def transform_mnist_data(x, transform_mode, max_translation=5, sigma=0.1, show=False):
"""
Transform MNIST data to generate appropriate data for transforming autoencoder training.
Parameters
----------
x: ndarray
MNIST images
transform_mode: str
Transformation in ['translation', 'affine']
max_translation: int
Maximum translation allowed in the transformation matrix
sigma: float
Parametrizes the gaussian noise added to transformation matrix
show: bool
If True both original and transformed images are shown
Returns
-------
data: dict
Dictionary containing both original and transformed data
along with transformations applied.
"""
if transform_mode not in ['translation', 'affine']:
raise ValueError('Mode "{}" not supported.'.format(transform_mode))
if show:
plt.ion()
_, [ax1, ax2] = plt.subplots(1, 2)
examples = []
for i in np.random.permutation(len(x)): # notice shuffling here
mnist_image = np.reshape(x[i], (28, 28)) # MNIST is distributed flattened
if transform_mode == 'translation':
translation_x = random.randint(-max_translation, max_translation)
translation_y = random.randint(-max_translation, max_translation)
transformation = [translation_x, translation_y]
transformed_image = np.roll(np.roll(mnist_image, translation_x, axis=0), translation_y, axis=1)
else: # transform_mode == 'affine':
transformation = get_random_affine_matrix(sigma=sigma, max_translation=max_translation)
transformed_image = warp(mnist_image, AffineTransform(matrix=transformation))
if show:
ax1.imshow(mnist_image)
ax2.imshow(transformed_image)
plt.show()
plt.waitforbuttonpress()
example = TransformingAutoencoderExample(view_1=np.expand_dims(mnist_image, axis=-1),
view_2=np.expand_dims(transformed_image, axis=-1),
transformation=transformation)
examples.append(example)
return examples