-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmodels.py
124 lines (89 loc) · 3.89 KB
/
models.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
113
114
115
116
117
118
119
120
121
122
123
124
import tensorflow as tf
import tensorflow.contrib.slim as slim
def resnet(input_image):
'''
https://github.com/jmiller656/EDSR-Tensorflow/blob/master/model.py
EDSR paper: num_layers=32,feature_size=256
dped paper: num_layers=4,feature_size=64
'''
with tf.variable_scope("generator"):
#Preprocessing as mentioned in the EDSR paper, by subtracting the mean of the entire dataset. Here subtracting the mean of each batch.
mean_input_image=tf.reduce_mean(input_image)
image_input=input_image-mean_input_image
num_layers=32
feature_size=64
output_channels=3
w= weight_variable([9, 9, 36, feature_size]); b = bias_variable([feature_size]);
x = conv2d(input_image,w )+b
# Store the output of the first convolution to add later
conv_1 = x
for i in range(num_layers):
# x = ResBlock_NoBN(x)
x = ResBlock(x)
#three more convolution, and then we add the output of our first conv layer
w= weight_variable([3, 3, feature_size, feature_size]); b = bias_variable([feature_size]);
x = conv2d(x, w)+b
w= weight_variable([3, 3, feature_size, feature_size]); b = bias_variable([feature_size]);
x = conv2d(x, w)+b
x += conv_1
w= weight_variable([9, 9, feature_size, 3]); b = bias_variable([3]);
x = conv2d(x, w)+b
# Final
# enhanced = tf.clip_by_value(x,0.0,1.0)
enhanced = tf.nn.tanh(x) * 0.5 + 0.5
return enhanced
"""
Creates a convolutional residual block as defined in the paper. More on this inside model.py
x: input to pass through the residual block
channels: number of channels to compute
stride: convolution stride
"""
def ResBlock(x):
w= weight_variable([3, 3, 64, 64]); b = bias_variable([64])
tmp = tf.nn.relu(_instance_norm(conv2d(x, w)+b))
w= weight_variable([3, 3, 64, 64]); b = bias_variable([64])
tmp = tf.nn.relu(_instance_norm(conv2d(tmp, w)+b))
return x + tmp
def ResBlock_NoBN(x,channels=64,kernel_size=[3,3]):
scale=0.1
tmp = slim.conv2d(x,channels,kernel_size,activation_fn=None)
tmp = tf.nn.relu(tmp)
tmp = slim.conv2d(tmp,channels,kernel_size,activation_fn=None)
tmp*=scale
return x + tmp
def weight_variable(shape, name='w'):
initial = tf.truncated_normal(shape, stddev=0.01)
return tf.Variable(initial, name=name)
def bias_variable(shape, name='b'):
initial = tf.constant(0.01, shape=shape)
return tf.Variable(initial, name=name)
def conv2d(x, W):
return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')
def leaky_relu(x, alpha = 0.2):
return tf.maximum(alpha * x, x)
def _conv_layer(net, num_filters, filter_size, strides, batch_nn=True):
weights_init = _conv_init_vars(net, num_filters, filter_size)
strides_shape = [1, strides, strides, 1]
bias = tf.Variable(tf.constant(0.01, shape=[num_filters]))
net = tf.nn.conv2d(net, weights_init, strides_shape, padding='SAME') + bias
net = leaky_relu(net)
if batch_nn:
net = _instance_norm(net)
return net
def _instance_norm(net):
batch, rows, cols, channels = [i.value for i in net.get_shape()]
var_shape = [channels]
mu, sigma_sq = tf.nn.moments(net, [1,2], keep_dims=True)
shift = tf.Variable(tf.zeros(var_shape))
scale = tf.Variable(tf.ones(var_shape))
epsilon = 1e-3
normalized = (net-mu)/(sigma_sq + epsilon)**(.5)
return scale * normalized + shift
def _conv_init_vars(net, out_channels, filter_size, transpose=False):
_, rows, cols, in_channels = [i.value for i in net.get_shape()]
if not transpose:
weights_shape = [filter_size, filter_size, in_channels, out_channels]
else:
weights_shape = [filter_size, filter_size, out_channels, in_channels]
weights_init = tf.Variable(tf.truncated_normal(weights_shape, stddev=0.01, seed=1), dtype=tf.float32)
return weights_init