-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathCAE.py
98 lines (76 loc) · 3.83 KB
/
CAE.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
# -*- coding: utf-8 -*-
"""
@author: Shahin Heidarian
Concordia Intelligent Signal & Information Processing (I-SIP)
"""
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import *
from tensorflow.keras import Model
def build_encoder(input_shape = (512,512,1), conv_dims = [32,64,128], dense_dims = [1024], activation = 'relu', show_summary = True):
encoder_inputs = Input(shape=input_shape)
x = Conv2D(conv_dims[0], 3, activation=activation, strides=1, padding="same", trainable = True)(encoder_inputs)
x = MaxPooling2D((2, 2), padding="same")(x)
if len(conv_dims)>1:
for dim in conv_dims[1::]:
x = Conv2D(dim, 3, activation=activation, strides=1, padding="same", trainable = True)(x)
x = MaxPooling2D((2, 2), padding="same")(x)
x = Flatten()(x)
x = Dense(dense_dims[0], activation=activation, trainable = True)(x)
if len(dense_dims) >1:
for dens in dense_dims[1::]:
x = Dense(dens, activation=activation, trainable = True)(x)
encoder = Model(inputs = encoder_inputs, outputs = x, name="encoder")
if show_summary:
encoder.summary()
return encoder
def build_encoder_with_loc(input_shape = (512,512,1), conv_dims = [32,64,128], dense_dims = [1024], activation = 'relu', show_summary = True):
encoder_inputs = Input(shape=input_shape)
x = Conv2D(conv_dims[0], 3, activation=activation, strides=1, padding="same", trainable = True)(encoder_inputs)
x = MaxPooling2D((2, 2), padding="same")(x)
if len(conv_dims)>1:
for dim in conv_dims[1::]:
x = Conv2D(dim, 3, activation=activation, strides=1, padding="same", trainable = True)(x)
x = MaxPooling2D((2, 2), padding="same")(x)
x = Flatten()(x)
x = Dense(dense_dims[0], activation=activation, trainable = True)(x)
if len(dense_dims) >1:
for dens in dense_dims[1::]:
x = Dense(dens, activation=activation, trainable = True)(x)
encoder = Model(inputs = encoder_inputs, outputs = x, name="encoder")
if show_summary:
encoder.summary()
return encoder
def build_decoder(output_size = 512, deconv_dims = [128,64,32], dense_dims = [1024], activation = 'relu', show_summary = True):
latent_inputs = Input(shape=(dense_dims[0]))
num_deconv = len(deconv_dims)
dense_neurons = int(output_size/(2**num_deconv)) # it is a quotient of the output size, so that the image will be reconstructed to the right size correctly
x = Dense(dense_neurons * dense_neurons * dense_dims[0], activation=activation, trainable = True)(latent_inputs)
x = Reshape((dense_neurons, dense_neurons, dense_dims[0]))(x)
x = Conv2DTranspose(deconv_dims[0], 3, activation=activation, strides=2, padding="same", trainable = True)(x)
if num_deconv >1:
for dim in deconv_dims[1::]:
x = Conv2DTranspose(dim, 3, activation=activation, strides=2, padding="same", trainable = True)(x)
decoder_outputs = Conv2DTranspose(1, 3, activation="sigmoid", padding="same", trainable = True)(x)
decoder = Model(latent_inputs, decoder_outputs, name="decoder")
if show_summary:
decoder.summary()
return decoder
class CAE(Model):
def __init__(self,
encoder,
decoder,
model_input_shape = (512,512,1),
**kwargs):
super(CAE, self).__init__(**kwargs)
self.encoder = encoder
self.decoder = decoder
self.model_input_shape = model_input_shape
def call(self, inputs):
encoder_outputs = self.encoder(inputs)
decoder_outputs = self.decoder(encoder_outputs)
return decoder_outputs
def model(self):
x = Input(shape = self.model_input_shape)
return Model(inputs=[x], outputs=self.call(x))