-
Notifications
You must be signed in to change notification settings - Fork 6
/
Copy pathVGG16AutoEncoderGenerator.py
89 lines (69 loc) · 4.06 KB
/
VGG16AutoEncoderGenerator.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
from keras.models import Sequential, Model
from keras.layers import Dense
from keras.layers.convolutional import Conv2D, MaxPooling2D, AveragePooling2D, UpSampling2D
class VGG16AutoEncoderGenerator:
def __init__(self, input_shape, weight_file: str, number_of_layers_to_freeze_from_input_layer: int):
self.encoder = VGG16EncoderGenerator(
input_shape=input_shape, weight_file=weight_file,
number_of_layers_to_freeze_from_input_layer=number_of_layers_to_freeze_from_input_layer
).encoder
self.decoder = DecoderGenerator(self.encoder).decoder
self.auto_encoder = Model(inputs=self.encoder.inputs, outputs=self.decoder)
class VGG16EncoderGenerator:
def __init__(self, input_shape, weight_file: str, number_of_layers_to_freeze_from_input_layer: int):
self.encoder = self.__get_encoder(input_shape)
self.encoder.load_weights(weight_file)
self.encoder = self.__freeze_weight(self.encoder, number_of_layers_to_freeze_from_input_layer)
def __get_encoder(self, input_shape):
return Sequential([
Conv2D(64, (3, 3), input_shape=input_shape, padding='same', activation='relu'),
Conv2D(64, (3, 3), activation='relu', padding='same'),
MaxPooling2D(pool_size=(2, 2), strides=(2, 2)),
Conv2D(128, (3, 3), activation='relu', padding='same'),
Conv2D(128, (3, 3), activation='relu', padding='same',),
MaxPooling2D(pool_size=(2, 2), strides=(2, 2)),
Conv2D(256, (3, 3), activation='relu', padding='same',),
Conv2D(256, (3, 3), activation='relu', padding='same',),
Conv2D(256, (3, 3), activation='relu', padding='same',),
MaxPooling2D(pool_size=(2, 2), strides=(2, 2)),
Conv2D(512, (3, 3), activation='relu', padding='same',),
Conv2D(512, (3, 3), activation='relu', padding='same',),
Conv2D(512, (3, 3), activation='relu', padding='same',),
MaxPooling2D(pool_size=(2, 2), strides=(2, 2)),
Conv2D(512, (3, 3), activation='relu', padding='same',),
Conv2D(512, (3, 3), activation='relu', padding='same',),
Conv2D(512, (3, 3), activation='relu', padding='same',),
AveragePooling2D(pool_size=(2, 2), strides=(2, 2)),
])
def __freeze_weight(self, encoder, number_of_layers_to_freeze_from_input_layer: int):
for i in range(number_of_layers_to_freeze_from_input_layer):
encoder.layers[i].trainable = False
return encoder
class DecoderGenerator:
def __init__(self, encoder):
self.decoder = self.___get_decoder(encoder)
def ___get_decoder(self, encoder):
decoded_encoder = Dense(784, activation='relu')(encoder.output)
decoded_encoder = Conv2D(
filters=16, kernel_size=(3, 3), strides=(1, 1), padding='same', activation='relu'
)(decoded_encoder)
decoded_encoder = UpSampling2D(size=(2, 2))(decoded_encoder)
decoded_encoder = Conv2D(
filters=32, kernel_size=(3, 3), strides=(1, 1), padding='same', activation='relu'
)(decoded_encoder)
decoded_encoder = UpSampling2D(size=(2, 2))(decoded_encoder)
decoded_encoder = Conv2D(
filters=64, kernel_size=(3, 3), strides=(1, 1), padding='same', activation='relu'
)(decoded_encoder)
decoded_encoder = UpSampling2D(size=(2, 2))(decoded_encoder)
decoded_encoder = Conv2D(
filters=128, kernel_size=(3, 3), strides=(1, 1), padding='same', activation='relu'
)(decoded_encoder)
decoded_encoder = UpSampling2D(size=(2, 2))(decoded_encoder)
decoded_encoder = Conv2D(
filters=64, kernel_size=(3, 3), strides=(1, 1), padding='same', activation='relu')(decoded_encoder)
decoded_encoder = UpSampling2D(size=(2, 2))(decoded_encoder)
decoder = Conv2D(
filters=3, kernel_size=(3, 3), strides=(1, 1), padding='same', activation='sigmoid'
)(decoded_encoder)
return decoder