=
61 lines
2.9 KiB
Python
61 lines
2.9 KiB
Python
from tensorflow.keras.layers import Conv2D, UpSampling2D, LeakyReLU, Concatenate
|
|
from tensorflow.keras import Model
|
|
from tensorflow.keras.applications import DenseNet169
|
|
|
|
class UpscaleBlock(Model):
|
|
def __init__(self, filters, name):
|
|
super(UpscaleBlock, self).__init__()
|
|
self.up = UpSampling2D(size=(2, 2), interpolation='bilinear', name=name+'_upsampling2d')
|
|
self.concat = Concatenate(name=name+'_concat') # Skip connection
|
|
self.convA = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same', name=name+'_convA')
|
|
self.reluA = LeakyReLU(alpha=0.2)
|
|
self.convB = Conv2D(filters=filters, kernel_size=3, strides=1, padding='same', name=name+'_convB')
|
|
self.reluB = LeakyReLU(alpha=0.2)
|
|
|
|
def call(self, x):
|
|
b = self.reluB( self.convB( self.reluA( self.convA( self.concat( [self.up(x[0]), x[1]] ) ) ) ) )
|
|
return b
|
|
|
|
class Encoder(Model):
|
|
def __init__(self):
|
|
super(Encoder, self).__init__()
|
|
self.base_model = DenseNet169(input_shape=(None, None, 3), include_top=False, weights='imagenet')
|
|
print('Base model loaded {}'.format(DenseNet169.__name__))
|
|
|
|
# Create encoder model that produce final features along with multiple intermediate features
|
|
outputs = [self.base_model.outputs[-1]]
|
|
for name in ['pool1', 'pool2_pool', 'pool3_pool', 'conv1/relu'] : outputs.append( self.base_model.get_layer(name).output )
|
|
self.encoder = Model(inputs=self.base_model.inputs, outputs=outputs)
|
|
|
|
def call(self, x):
|
|
return self.encoder(x)
|
|
|
|
class Decoder(Model):
|
|
def __init__(self, decode_filters):
|
|
super(Decoder, self).__init__()
|
|
self.conv2 = Conv2D(filters=decode_filters, kernel_size=1, padding='same', name='conv2')
|
|
self.up1 = UpscaleBlock(filters=decode_filters//2, name='up1')
|
|
self.up2 = UpscaleBlock(filters=decode_filters//4, name='up2')
|
|
self.up3 = UpscaleBlock(filters=decode_filters//8, name='up3')
|
|
self.up4 = UpscaleBlock(filters=decode_filters//16, name='up4')
|
|
self.conv3 = Conv2D(filters=1, kernel_size=3, strides=1, padding='same', name='conv3')
|
|
|
|
def call(self, features):
|
|
x, pool1, pool2, pool3, conv1 = features[0], features[1], features[2], features[3], features[4]
|
|
up0 = self.conv2(x)
|
|
up1 = self.up1([up0, pool3])
|
|
up2 = self.up2([up1, pool2])
|
|
up3 = self.up3([up2, pool1])
|
|
up4 = self.up4([up3, conv1])
|
|
return self.conv3( up4 )
|
|
|
|
class DepthEstimate(Model):
|
|
def __init__(self):
|
|
super(DepthEstimate, self).__init__()
|
|
self.encoder = Encoder()
|
|
self.decoder = Decoder( decode_filters = int(self.encoder.layers[-1].output[0].shape[-1] // 2 ) )
|
|
print('\nModel created.')
|
|
|
|
def call(self, x):
|
|
return self.decoder( self.encoder(x) )
|