def build_model(self, input_shape, cnum=32):
"""
Build up the PEN-Net structure(Discriminator,Generator and Combined Model)
:param input_shape: Input Shape for the PEN-Net (channel[0:2] for Img,channel[-1] for mask)
:param cnum: Conv channel number
:return: Discriminator,Generator and Combined Model
"""
img_input = Input(batch_shape=input_shape[0])
mask_input = Input(batch_shape=input_shape[1])
img_shape = [input_shape[0][0], input_shape[0][1], input_shape[0][2], input_shape[0][3] - 1]
# encoder
dw_conv01 = Sequential([Conv2D(cnum, kernel_size=3, strides=2, padding='same'),
LeakyReLU(0.2)])(img_input)
dw_conv02 = Sequential([Conv2D(cnum * 2, kernel_size=3, strides=2, padding='same'),
(LeakyReLU(0.2))])(dw_conv01)
dw_conv03 = Sequential([Conv2D(cnum * 4, kernel_size=3, strides=2, padding='same'),
(LeakyReLU(0.2))])(dw_conv02)
dw_conv04 = Sequential([Conv2D(cnum * 8, kernel_size=3, strides=2, padding='same'),
(LeakyReLU(0.2))])(dw_conv03)
dw_conv05 = Sequential([Conv2D(cnum * 16, kernel_size=3, strides=2, padding='same'),
(LeakyReLU(0.2))])(dw_conv04)
dw_conv06 = Sequential([Conv2D(cnum * 16, kernel_size=3, strides=2, padding='same'),
(LeakyReLU(0.2))])(dw_conv05)
dw_conv06_shape = dw_conv06.get_shape().as_list()
dw_conv05_shape = dw_conv05.get_shape().as_list()
dw_conv04_shape = dw_conv04.get_shape().as_list()
dw_conv03_shape = dw_conv03.get_shape().as_list()
dw_conv02_shape = dw_conv02.get_shape().as_list()
dw_conv01_shape = dw_conv01.get_shape().as_list()
mask_shape = mask_input.get_shape().as_list()
at_conv05 = ATNConv(dw_conv05_shape, dw_conv06_shape, mask_shape, 1, False)(
[dw_conv05, dw_conv06, mask_input])
at_conv04 = ATNConv(dw_conv04_shape, dw_conv05_shape, mask_shape)([dw_conv04, dw_conv05, mask_input])
at_conv03 = ATNConv(dw_conv03_shape, dw_conv04_shape, mask_shape)([dw_conv03, dw_conv04, mask_input])
at_conv02 = ATNConv(dw_conv02_shape, dw_conv03_shape, mask_shape)([dw_conv02, dw_conv03, mask_input])
at_conv01 = ATNConv(dw_conv01_shape, dw_conv02_shape, mask_shape)([dw_conv01, dw_conv02, mask_input])
upx5_ = UpSampling2D((2, 2), interpolation='bilinear')(dw_conv06)
up_conv05 = Conv2D(cnum * 16, kernel_size=3, strides=1, padding="same", activation='relu')(upx5_)
upx4_ = UpSampling2D((2, 2), interpolation='bilinear')(Concatenate(axis=-1)([up_conv05, at_conv05]))
up_conv04 = Conv2D(cnum * 8, kernel_size=3, strides=1, padding="same", activation='relu')(upx4_)
upx3_ = UpSampling2D((2, 2), interpolation='bilinear')(Concatenate(axis=-1)([up_conv04, at_conv04]))
up_conv03 = Conv2D(cnum * 4, kernel_size=3, strides=1, padding="same", activation='relu')(upx3_)
upx2_ = UpSampling2D((2, 2), interpolation='bilinear')(Concatenate(axis=-1)([up_conv03, at_conv03]))
up_conv02 = Conv2D(cnum * 2, kernel_size=3, strides=1, padding="same", activation='relu')(upx2_)
upx1_ = UpSampling2D((2, 2), interpolation='bilinear')(Concatenate(axis=-1)([up_conv02, at_conv02]))
up_conv01 = Conv2D(cnum, kernel_size=3, strides=1, padding="same", activation='relu')(upx1_)
torgb5 = Conv2D(3, kernel_size=1, strides=1, padding="same", activation='tanh')
torgb4 = Conv2D(3, kernel_size=1, strides=1, padding="same", activation='tanh')
torgb3 = Conv2D(3, kernel_size=1, strides=1, padding="same", activation='tanh')
torgb2 = Conv2D(3, kernel_size=1, strides=1, padding="same", activation='tanh')
torgb1 = Conv2D(3, kernel_size=1, strides=1, padding="same", activation='tanh')
img5 = torgb5(Concatenate(axis=-1)([up_conv05, at_conv05]))
img4 = torgb4(Concatenate(axis=-1)([up_conv04, at_conv04]))
img3 = torgb3(Concatenate(axis=-1)([up_conv03, at_conv03]))
img2 = torgb2(Concatenate(axis=-1)([up_conv02, at_conv02]))
img1 = torgb1(Concatenate(axis=-1)([up_conv01, at_conv01]))
output_ = UpSampling2D((2, 2), interpolation='bilinear')(Concatenate(axis=-1)([up_conv01, at_conv01]))
decoded_out = Sequential([
Conv2D(cnum * 2, kernel_size=3, strides=1, padding="same", activation='relu'),
Conv2D(3, kernel_size=3, strides=1, padding="same", activation='tanh')])(output_)
mask_out = Lambda(Mask_Img_layer)([decoded_out, mask_input])
unmask_out = Lambda(UnMasked_Img_layer)([decoded_out, mask_input])
Discriminator = self.build_discriminator(img_shape)
Discriminator.compile(loss='binary_crossentropy',
optimizer=self.optimizer,
metrics=['accuracy'])
Discriminator.trainable = False # Combined Model's Discriminator can not be trained
discrim_out = Discriminator(decoded_out)
Generator = Model([img_input, mask_input],
[decoded_out, mask_out, unmask_out, img1, img2, img3, img4, img5], name="Generator")
Combined_Model = Model([img_input, mask_input],
[mask_out, unmask_out, img1, img2, img3, img4, img5, discrim_out],
name="Combined_Model")
print("Generator Structure:")
Generator.summary()
print("Combined_Model Structure:")
Combined_Model.summary()
return Combined_Model, Generator, Discriminator
