#!/usr/bin/python3
from tensorflow import Tensor as tf_tensor
from tensorflow.python.keras.layers import (
Activation,
BatchNormalization,
Conv2D,
Input,
MaxPooling2D,
SpatialDropout2D,
UpSampling2D,
concatenate,
)
from tensorflow.python.keras.models import Model
[docs]class ModelLayers:
"""Util Class to hold all methods to create model layers for UNAGI Model."""
[docs] @staticmethod
def double_conv_layer(inputs: tf_tensor, kernel_size: int) -> tf_tensor:
"""Create a double convolution layer with mentioned inputs and filters.
Parameters
----------
inputs: tf_tensor
input keras layer
kernel_size: int
filter size/ kernel size for the tf_tensor
Returns
-------
tf_tensor
tensor with mentioned number of filters
Example
-------
unagi.utils.model_utils.ModelLayers.double_conv_layer(Input, 32)
"""
conv = Conv2D(
kernel_size, (3, 3), padding="same", kernel_initializer="he_normal"
)(inputs)
conv = BatchNormalization(axis=3)(conv)
conv = Activation("relu")(conv)
conv = Conv2D(
kernel_size, (3, 3), padding="same", kernel_initializer="he_normal"
)(conv)
conv = BatchNormalization(axis=3)(conv)
conv = Activation("relu")(conv)
conv = SpatialDropout2D(0.1)(conv)
return conv
[docs] @staticmethod
def down_layer(inputs: tf_tensor, kernel_size: int) -> tf_tensor:
"""Create downsampling layer.
Parameters
----------
inputs: tf_tensor
input keras layer
kernel_size: int
filter size/ kernel size for the tf_tensor
Returns
-------
tf_tensor
tensor with mentioned number of filters
See Also
--------
ModelLayers.double_conv_layer()
Example
-------
unagi.utils.model_utils.ModelLayers.down_layer(tf_tensor, 64)
"""
conv = ModelLayers.double_conv_layer(inputs, kernel_size)
pool = MaxPooling2D(pool_size=(2, 2))(conv)
return conv, pool
[docs] @staticmethod
def up_layer(inputs: tf_tensor, concats: tf_tensor, kernel_size: int) -> tf_tensor:
"""Create upsampling layer.
Parameters
----------
inputs: tf_tensor
input keras layer
concats: tf_tensor
keras layer to concat
kernel_size: int
filter size/ kernel size for the tf_tensor
Returns
-------
tf_tensor
tensor with mentioned number of filters
See Also
--------
ModelLayers.double_conv_layer()
Example
-------
unagi.utils.model_utils.ModelLayers.up_layer(input, pool_layer, 128)
"""
return ModelLayers.double_conv_layer(
concatenate([UpSampling2D(size=(2, 2))(inputs), concats], axis=3),
kernel_size,
)
[docs]class UNAGIModel:
"""Util Class to create U-net model for UNAGI."""
def __init__(self, width: int = 128, height: int = 128, channels: int = 1) -> None:
"""Create a U-net model class.
Parameters
----------
width : int
width of input tensor
height : int
height of input tensor
channels : int
number of channels in image
"""
self.width = width
self.height = height
self.channels = channels
[docs] def unet(self) -> Model:
"""Create U-net for input image/tensor size.
Returns
-------
Model
keras Model with mentioned number of filters and layers
See Also
--------
ModelLayers.double_conv_layer(), ModelLayers.down_layer(),
ModelLayers.up_layer()
Example
-------
unagi.utils.model_utils.UNetModel.unet()
"""
inputs = Input((self.width, self.height, self.channels))
model_layers = ModelLayers()
# Downsampling.
down1, pool1 = model_layers.down_layer(inputs, 32)
down2, pool2 = model_layers.down_layer(pool1, 64)
down3, pool3 = model_layers.down_layer(pool2, 128)
down4, pool4 = model_layers.down_layer(pool3, 256)
down5, pool5 = model_layers.down_layer(pool4, 512)
# Bottleneck.
bottleneck = model_layers.double_conv_layer(pool5, 1024)
# Upsampling.
up5 = model_layers.up_layer(bottleneck, down5, 512)
up4 = model_layers.up_layer(up5, down4, 256)
up3 = model_layers.up_layer(up4, down3, 128)
up2 = model_layers.up_layer(up3, down2, 64)
up1 = model_layers.up_layer(up2, down1, 32)
outputs = Conv2D(1, (1, 1))(up1)
outputs = Activation("sigmoid")(outputs)
model = Model(inputs, outputs)
return model