felixrosberg commited on
Commit
3b88fc9
β€’
1 Parent(s): fae5def

removed generator code

Browse files
networks/__pycache__/generator.cpython-37.pyc DELETED
Binary file (7.03 kB)
 
networks/__pycache__/generator.cpython-38.pyc DELETED
Binary file (6.54 kB)
 
networks/generator.py DELETED
@@ -1,321 +0,0 @@
1
- from tensorflow.keras.layers import *
2
- from tensorflow.keras.models import Model
3
- from tensorflow_addons.layers import InstanceNormalization
4
- from networks.layers import AdaIN, AdaptiveAttention
5
-
6
- import numpy as np
7
-
8
-
9
- def residual_down_block(inputs, filters, resample=True):
10
- x = inputs
11
-
12
- r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
13
- if resample:
14
- r = AveragePooling2D()(r)
15
-
16
- x = InstanceNormalization()(x)
17
- x = LeakyReLU(0.2)(x)
18
- x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
19
-
20
- if resample:
21
- x = AveragePooling2D()(x)
22
-
23
- x = Add()([x, r])
24
-
25
- return x
26
-
27
-
28
- def residual_up_block(inputs, filters, resample=True, name=None):
29
- x, z_id = inputs
30
-
31
- r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
32
- if resample:
33
- r = UpSampling2D(interpolation='bilinear')(r)
34
-
35
- x = InstanceNormalization()(x)
36
- x = AdaIN()([x, z_id])
37
- x = LeakyReLU(0.2)(x)
38
- x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
39
-
40
- if resample:
41
- x = UpSampling2D(interpolation='bilinear')(x)
42
-
43
- x = Add()([x, r])
44
-
45
- return x
46
-
47
-
48
- def adaptive_attention(inputs, filters, name=None):
49
- x_t, x_s = inputs
50
-
51
- m = Concatenate(axis=-1)([x_t, x_s])
52
- m = Conv2D(filters=filters // 4, kernel_size=3, strides=1, padding='same')(m)
53
- m = LeakyReLU(0.2)(m)
54
- m = InstanceNormalization()(m)
55
- m = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same', activation='sigmoid', name=name)(m)
56
-
57
- x = AdaptiveAttention()([m, x_t, x_s])
58
-
59
- return x
60
-
61
-
62
- def adaptive_fusion_up_block(inputs, filters, resample=True, name=None):
63
- x_t, x_s, z_id = inputs
64
-
65
- x = adaptive_attention([x_t, x_s], x_t.shape[-1], name=name)
66
-
67
- r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
68
- if resample:
69
- r = UpSampling2D(interpolation='bilinear')(r)
70
-
71
- x = InstanceNormalization()(x)
72
- x = AdaIN()([x, z_id])
73
- x = LeakyReLU(0.2)(x)
74
- x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
75
-
76
- if resample:
77
- x = UpSampling2D(interpolation='bilinear')(x)
78
-
79
- x = Add()([x, r])
80
-
81
- return x
82
-
83
-
84
- def dual_adaptive_fusion_up_block(inputs, filters, resample=True, name=None):
85
- x_t, x_s, z_id = inputs
86
-
87
- x = adaptive_attention([x_t, x_s], x_t.shape[-1], name=name + '_0')
88
- x = adaptive_attention([x_t, x], x_t.shape[-1], name=name + '_1')
89
-
90
- r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
91
- if resample:
92
- r = UpSampling2D(interpolation='bilinear')(r)
93
-
94
- x = InstanceNormalization()(x)
95
- x = AdaIN()([x, z_id])
96
- x = LeakyReLU(0.2)(x)
97
- x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
98
-
99
- if resample:
100
- x = UpSampling2D(interpolation='bilinear')(x)
101
-
102
- x = Add()([x, r])
103
-
104
- return x
105
-
106
-
107
- def adaptive_fusion_up_block_no_add(inputs, filters, resample=True, name=None):
108
- x_t, x_s, z_id = inputs
109
-
110
- x = adaptive_attention([x_t, x_s], x_t.shape[-1], name=name)
111
-
112
- x = InstanceNormalization()(x)
113
- x = AdaIN()([x, z_id])
114
- x = LeakyReLU(0.2)(x)
115
- x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
116
-
117
- return x
118
-
119
-
120
- def adaptive_fusion_up_block_concat_baseline(inputs, filters, resample=True, name=None):
121
- x_t, x_s, z_id = inputs
122
-
123
- x = Concatenate(axis=-1)([x_t, x_s])
124
-
125
- r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
126
- if resample:
127
- r = UpSampling2D(interpolation='bilinear')(r)
128
-
129
- x = InstanceNormalization()(x)
130
- x = AdaIN()([x, z_id])
131
- x = LeakyReLU(0.2)(x)
132
- x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
133
-
134
- if resample:
135
- x = UpSampling2D(interpolation='bilinear')(x)
136
-
137
- x = Add(name=name if name == 'final' else None)([x, r])
138
-
139
- return x
140
-
141
-
142
- def adaptive_fusion_up_block_add_baseline(inputs, filters, resample=True, name=None):
143
- x_t, x_s, z_id = inputs
144
-
145
- x = Add()([x_t, x_s])
146
-
147
- r = Conv2D(filters=filters, kernel_size=1, strides=1, padding='same')(x)
148
- if resample:
149
- r = UpSampling2D(interpolation='bilinear')(r)
150
-
151
- x = InstanceNormalization()(x)
152
- x = AdaIN()([x, z_id])
153
- x = LeakyReLU(0.2)(x)
154
- x = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same')(x)
155
-
156
- if resample:
157
- x = UpSampling2D(interpolation='bilinear')(x)
158
-
159
- x = Add()([x, r])
160
-
161
- return x
162
-
163
-
164
- def get_generator_original(mapping_depth=4, mapping_size=256):
165
- x_target = Input(shape=(256, 256, 3))
166
- z_source = Input(shape=(512,))
167
-
168
- z_id = z_source
169
- for m in range(np.max([mapping_depth - 1, 0])):
170
- z_id = Dense(mapping_size)(z_id)
171
- z_id = LeakyReLU(0.2)(z_id)
172
- if mapping_depth >= 1:
173
- z_id = Dense(mapping_size)(z_id)
174
-
175
- x_0 = Conv2D(filters=64, kernel_size=3, strides=1, padding='same')(x_target) # 256
176
-
177
- x_1 = residual_down_block(x_0, 128) # 128
178
-
179
- x_2 = residual_down_block(x_1, 256) # 64
180
-
181
- x_3 = residual_down_block(x_2, 512)
182
-
183
- x_4 = residual_down_block(x_3, 512)
184
-
185
- x_5 = residual_down_block(x_4, 512)
186
-
187
- x_6 = residual_down_block(x_5, 512, resample=False)
188
-
189
- u_5 = residual_up_block([x_6, z_id], 512, resample=False)
190
-
191
- u_4 = residual_up_block([u_5, z_id], 512)
192
-
193
- u_3 = residual_up_block([u_4, z_id], 512)
194
-
195
- u_2 = residual_up_block([u_3, z_id], 256) # 64
196
-
197
- u_1 = adaptive_fusion_up_block([x_2, u_2, z_id], 128, name='aff_attention_64x64') # 128
198
-
199
- u_0 = adaptive_fusion_up_block([x_1, u_1, z_id], 64, name='aff_attention_128x128') # 256
200
-
201
- out = adaptive_fusion_up_block([x_0, u_0, z_id], 3, resample=False, name='aff_attention_256x256')
202
-
203
- gen_model = Model([x_target, z_source], out)
204
- gen_model.summary()
205
-
206
- return gen_model
207
-
208
-
209
- def make_layer(l_type, inputs, filters, resample, name=None):
210
- if l_type == 'affa':
211
- return adaptive_fusion_up_block(inputs, filters, resample=resample, name=name)
212
- if l_type == 'd_affa':
213
- return dual_adaptive_fusion_up_block(inputs, filters, resample=resample, name=name)
214
- elif l_type == 'concat':
215
- return adaptive_fusion_up_block_concat_baseline(inputs, filters, resample=resample, name=name)
216
- elif l_type == 'no_skip':
217
- return residual_up_block(inputs[1:], filters, resample=resample)
218
-
219
-
220
- def get_generator(up_types=None, mapping_depth=4, mapping_size=256):
221
-
222
- if up_types is None:
223
- up_types = ['no_skip', 'no_skip', 'd_affa', 'd_affa', 'd_affa', 'concat']
224
-
225
- x_target = Input(shape=(256, 256, 3))
226
- z_source = Input(shape=(512,))
227
-
228
- z_id = z_source
229
- for m in range(np.max([mapping_depth - 1, 0])):
230
- z_id = Dense(mapping_size)(z_id)
231
- z_id = LeakyReLU(0.2)(z_id)
232
- if mapping_depth >= 1:
233
- z_id = Dense(mapping_size)(z_id)
234
-
235
- x_0 = Conv2D(filters=64, kernel_size=3, strides=1, padding='same')(x_target) # 256
236
-
237
- x_1 = residual_down_block(x_0, 128) # 128
238
-
239
- x_2 = residual_down_block(x_1, 256) # 64
240
-
241
- x_3 = residual_down_block(x_2, 512)
242
-
243
- x_4 = residual_down_block(x_3, 512)
244
-
245
- x_5 = residual_down_block(x_4, 512)
246
-
247
- x_6 = residual_down_block(x_5, 512, resample=False)
248
-
249
- u_5 = residual_up_block([x_6, z_id], 512, resample=False)
250
-
251
- u_4 = make_layer(up_types[0], [x_5, u_5, z_id], 512, resample=True, name='16x16')
252
-
253
- u_3 = make_layer(up_types[1], [x_4, u_4, z_id], 512, resample=True, name='32x32')
254
-
255
- u_2 = make_layer(up_types[2], [x_3, u_3, z_id], 256, resample=True, name='64x64')
256
-
257
- u_1 = make_layer(up_types[3], [x_2, u_2, z_id], 128, resample=True, name='128x128')
258
-
259
- u_0 = make_layer(up_types[4], [x_1, u_1, z_id], 64, resample=True, name='256x256')
260
-
261
- out = make_layer(up_types[5], [x_0, u_0, z_id], 3, resample=False, name='final')
262
-
263
- gen_model = Model([x_target, z_source], out)
264
- gen_model.summary()
265
-
266
- return gen_model
267
-
268
-
269
- def get_generator_large(up_types=None, mapping_depth=4, mapping_size=512):
270
-
271
- if up_types is None:
272
- up_types = ['no_skip', 'no_skip', 'affa', 'affa', 'affa', 'concat']
273
-
274
- x_target = Input(shape=(256, 256, 3))
275
- z_source = Input(shape=(512,))
276
-
277
- z_id = z_source
278
- for m in range(np.max([mapping_depth - 1, 0])):
279
- z_id = Dense(mapping_size)(z_id)
280
- z_id = LeakyReLU(0.2)(z_id)
281
- if mapping_depth >= 1:
282
- z_id = Dense(mapping_size)(z_id)
283
-
284
- x_0 = Conv2D(filters=64, kernel_size=3, strides=1, padding='same')(x_target) # 256
285
-
286
- x_1 = residual_down_block(x_0, 128) # 128
287
-
288
- x_2 = residual_down_block(x_1, 256) # 64
289
-
290
- x_3 = residual_down_block(x_2, 512)
291
-
292
- x_4 = residual_down_block(x_3, 512)
293
-
294
- x_5 = residual_down_block(x_4, 512)
295
-
296
- b_0 = residual_up_block([x_5, z_id], 512, resample=False)
297
-
298
- b_1 = residual_up_block([b_0, z_id], 512, resample=False)
299
-
300
- b_2 = residual_up_block([b_1, z_id], 512, resample=False)
301
-
302
- u_5 = residual_up_block([b_2, z_id], 512, resample=False)
303
-
304
- u_4 = make_layer(up_types[0], [x_5, u_5, z_id], 512, resample=True, name='16x16')
305
-
306
- u_3 = make_layer(up_types[1], [x_4, u_4, z_id], 512, resample=True, name='32x32')
307
-
308
- u_2 = make_layer(up_types[2], [x_3, u_3, z_id], 256, resample=True, name='64x64')
309
-
310
- u_1 = make_layer(up_types[3], [x_2, u_2, z_id], 128, resample=True, name='128x128')
311
-
312
- u_0 = make_layer(up_types[4], [x_1, u_1, z_id], 64, resample=True, name='256x256')
313
-
314
- out = make_layer(up_types[5], [x_0, u_0, z_id], 3, resample=False, name='final')
315
-
316
- gen_model = Model([x_target, z_source], out)
317
- gen_model.summary()
318
-
319
- return gen_model
320
-
321
-