models.py

"""
File:     deep-fus/src/models.py
Author:   Tommaso Di Ianni (todiian@stanford.edu)

Copyright 2021 Tommaso Di Ianni

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import numpy as np
import tensorflow as tf
# from tensorflow.keras import layers
from tensorflow.keras import regularizers
from tensorflow.keras.layers import Input, Conv2D, Conv3D, Activation, MaxPooling2D, Dropout, Concatenate, Conv2DTranspose, Reshape, Add
from tensorflow.keras.models import Model, load_model
from tensorflow.keras.initializers import he_uniform


def UNet5(input_shape = (96, 96, 125), do_rate=0.5):
    """
    Implementation of a 5-layer U-Net model for the reconstruction of power Doppler  images 
    from sparse compound data.
    
    Convolutional blocks are made of Conv2D + ReLU activations. 
    Downsampling is implemented with MaxPooling2D. 
    Upsampling is implemented with Conv2DTranspose. 
    Dropout used on all convolutional blocks.
    
    Arguments:
    input_shape -- dimensions of compounded dataset
    do_rate -- dropout factor   
    
    Returns:
    model -- a Model() instance in Keras
    
    """    
    
    def conv_block(X_in, nf, k, dr):
    
        X = Conv2D(filters=nf, kernel_size=(k,k), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X_in)
        X = Activation('relu')(X)
        X = Dropout(dr)(X)
    
        X = Conv2D(filters=nf, kernel_size=(k,k), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
        X = Activation('relu')(X)
        X = Dropout(dr)(X)
    
        return X
    
    # Filter kernel size used for convolutional layers
    filt_k = 3
    
    # Number of filters for respective layers
    F1 = 32
    F2 = 64
    F3 = 128
    F4 = 256
    F5 = 512
    
    # Input layer
    X_input = Input(input_shape, dtype=tf.float32, name="input")
    
    # Encoder L1
    X = conv_block(X_input, F1, filt_k, do_rate)
    
    X_skip1 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L2
    X = conv_block(X, F2, filt_k, do_rate)
    
    X_skip2 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L3
    X = conv_block(X, F3, filt_k, do_rate)
    
    X_skip3 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L4
    X = conv_block(X, F4, filt_k, do_rate)
    
    X_skip4 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L5
    X = conv_block(X, F5, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F4, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L4
    X = Concatenate(axis=3)([X,X_skip4])
   
    X = conv_block(X, F4, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F3, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L3
    X = Concatenate(axis=3)([X,X_skip3])
   
    X = conv_block(X, F3, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F2, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L2
    X = Concatenate(axis=3)([X,X_skip2])
   
    X = conv_block(X, F2, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F1, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L1    
    X = Concatenate(axis=3)([X,X_skip1])
   
    X = conv_block(X, F1, filt_k, do_rate)
    
    # Output layer
    X = Conv2D(filters=1, kernel_size=(1,1), strides=1, padding='same')(X)
    
    # Reshape output
    X = Reshape((input_shape[0],input_shape[1]))(X)
    
    # Create model
    model = Model(inputs = X_input, outputs = X, name='UNet5')
    
    return model



def UNet53D(input_shape = (96, 96, 125), do_rate=0.5, filt0_k12=3, filt0_k3=8):
    """
    Implementation of a 5-layer U-Net model for the reconstruction of power Doppler images 
    from sparse compound data. This model includes an initial Conv3D block that extracts 
    spatiotemporal features.
    
    Convolutional blocks are made of Conv2D + ReLU activations. 
    Downsampling is implemented with MaxPooling2D. 
    Upsampling is implemented with Conv2DTranspose. 
    Dropout used on all convolutional blocks.
    
    Arguments:
    input_shape -- dimensions of compound dataset
    do_rate -- dropout factor
    filt0_k12 -- first and second kernel dimensions for first Conv3D layer
    filt0_k3 -- third kernel dimension of first Conv3D layer
    
    Returns:
    model -- a Model() instance in Keras
    
    """
    
    def conv_block(X_in, nf, k, dr):
    
        X = Conv2D(filters=nf, kernel_size=(k,k), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X_in)
        X = Activation('relu')(X)
        X = Dropout(dr)(X)
    
        X = Conv2D(filters=nf, kernel_size=(k,k), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
        X = Activation('relu')(X)
        X = Dropout(dr)(X)
    
        return X
    
    # Filter kernel size used for convolutional layers
    filt_k = 3
    
    # Number of filters for respective layers
    F0 = 4
    
    F1 = 32
    F2 = 64
    F3 = 128
    F4 = 256
    F5 = 512
    
    # Input layer
    X_input = Input(input_shape, dtype=tf.float32, name="input")
    
    # Reshape input to 4-D
    X = Reshape((input_shape[0],input_shape[1],input_shape[2],1))(X_input)
    
    # Conv3D layer
    X = Conv3D(filters=F0, kernel_size=(filt0_k12,filt0_k12,filt0_k3), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
    X = Activation('relu')(X)
    
    # Reshape input to 3-D
    X = Reshape((input_shape[0],input_shape[1],input_shape[2]*F0))(X)
    
    # Encoder L1
    X = conv_block(X, F1, filt_k, do_rate)
    
    X_skip1 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L2
    X = conv_block(X, F2, filt_k, do_rate)
    
    X_skip2 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L3
    X = conv_block(X, F3, filt_k, do_rate)
    
    X_skip3 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L4
    X = conv_block(X, F4, filt_k, do_rate)
    
    X_skip4 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L5
    X = conv_block(X, F5, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F4, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L4
    X = Concatenate(axis=3)([X,X_skip4])
   
    X = conv_block(X, F4, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F3, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L3
    X = Concatenate(axis=3)([X,X_skip3])
   
    X = conv_block(X, F3, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F2, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L2
    X = Concatenate(axis=3)([X,X_skip2])
   
    X = conv_block(X, F2, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F1, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L1    
    X = Concatenate(axis=3)([X,X_skip1])
   
    X = conv_block(X, F1, filt_k, do_rate)
    
    # Output layer
    X = Conv2D(filters=1, kernel_size=(1,1), strides=1, padding='same')(X)
    
    # Reshape output
    X = Reshape((input_shape[0],input_shape[1]))(X)
    
    # Create model
    model = Model(inputs = X_input, outputs = X, name='UNet53D')
    
    return model



def ResUNet53D(input_shape = (96, 96, 125), do_rate=0.5, filt0_k12=3, filt0_k3=8):
    """
    Implementation of a 5-layer U-Net model with residual blocks for the reconstruction of power 
    Doppler images from sparse compound data. This model includes an initial Conv3D block that 
    extracts spatiotemporal features.
    
    Residual blocks are made of Conv2D + ReLU activations. 
    Downsampling is implemented with MaxPooling2D. 
    Upsampling is implemented with Conv2DTranspose. 
    Dropout used on all residual blocks.
    
    Arguments:
    input_shape -- dimensions of compound dataset
    do_rate -- dropout factor
    filt0_k12 -- first and second kernel dimensions for first Conv3D layer
    filt0_k3 -- third kernel dimension of first Conv3D layer
    
    Returns:
    model -- a Model() instance in Keras
    
    """
    
    def res_block(X_in, nf, k, dr):
    
        X = Conv2D(filters=nf, kernel_size=(k,k), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X_in)
        X = Activation('relu')(X)
        X = Dropout(dr)(X)
    
        X = Conv2D(filters=nf, kernel_size=(k,k), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
        X = Activation('relu')(X)
        X = Dropout(dr)(X)
        
        X_out = Add()([X, X_in])
    
        return X_out
    
    # Filter kernel size used for convolutional layers
    filt_k = 3
    
    # Number of filters for respective layers
    F0 = 4
    
    F1 = 32
    F2 = 64
    F3 = 128
    F4 = 256
    F5 = 512
    
    # Input layer
    X_input = Input(input_shape, dtype=tf.float32, name="input")
    
    # Reshape input to 4-D
    X = Reshape((input_shape[0],input_shape[1],input_shape[2],1))(X_input)
    
    # Conv3D layer
    X = Conv3D(filters=F0, kernel_size=(filt0_k12,filt0_k12,filt0_k3), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
    X = Activation('relu')(X)
    
    # Reshape input to 3-D
    X = Reshape((input_shape[0],input_shape[1],input_shape[2]*F0))(X)
    
    # Encoder L1
    X = Conv2D(filters=F1, kernel_size=(1,1), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
    
    X = res_block(X, F1, filt_k, do_rate)
    
    X_skip1 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L2
    X = Conv2D(filters=F2, kernel_size=(1,1), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
    
    X = res_block(X, F2, filt_k, do_rate)
    
    X_skip2 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L3
    X = Conv2D(filters=F3, kernel_size=(1,1), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
    
    X = res_block(X, F3, filt_k, do_rate)
    
    X_skip3 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L4
    X = Conv2D(filters=F4, kernel_size=(1,1), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
    
    X = res_block(X, F4, filt_k, do_rate)
    
    X_skip4 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L5
    X = Conv2D(filters=F5, kernel_size=(1,1), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
    
    X = res_block(X, F5, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F4, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L4
    X = Concatenate(axis=3)([X,X_skip4])
   
    X = Conv2D(filters=F4, kernel_size=(1,1), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
    
    X = res_block(X, F4, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F3, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L3
    X = Concatenate(axis=3)([X,X_skip3])
   
    X = Conv2D(filters=F3, kernel_size=(1,1), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
    
    X = res_block(X, F3, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F2, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L2
    X = Concatenate(axis=3)([X,X_skip2])
   
    X = Conv2D(filters=F2, kernel_size=(1,1), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
    
    X = res_block(X, F2, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F1, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L1    
    X = Concatenate(axis=3)([X,X_skip1])
   
    X = Conv2D(filters=F1, kernel_size=(1,1), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
    
    X = res_block(X, F1, filt_k, do_rate)
    
    # Output layer
    X = Conv2D(filters=1, kernel_size=(1,1), strides=1, padding='same')(X)
    
    # Reshape output
    X = Reshape((input_shape[0],input_shape[1]))(X)
    
    # Create model
    model = Model(inputs = X_input, outputs = X, name='ResUNet53D')
    
    return model



def UNet5_postproc(input_shape = (96, 96), do_rate=0.5):
    """
    Implementation of a 5-layer U-Net model for the post-processing of power Doppler images 
    reconstructed using sparse compound data.
    
    Convolutional blocks are made of Conv2D + ReLU activations. 
    Downsampling is implemented with MaxPooling2D. 
    Upsampling is implemented with Conv2DTranspose. 
    Dropout used on all convolutional blocks.
    
    Arguments:
    input_shape -- dimensions of compounded dataset
    do_rate -- dropout factor   
    
    Returns:
    model -- a Model() instance in Keras
    
    """    
    
    def conv_block(X_in, nf, k, dr):
    
        X = Conv2D(filters=nf, kernel_size=(k,k), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X_in)
        X = Activation('relu')(X)
        X = Dropout(dr)(X)
    
        X = Conv2D(filters=nf, kernel_size=(k,k), strides=1, padding='same', kernel_initializer=he_uniform(seed=0))(X)
        X = Activation('relu')(X)
        X = Dropout(dr)(X)
    
        return X
    
    # Filter kernel size used for convolutional layers
    filt_k = 3
    
    # Number of filters for respective layers
    F1 = 32
    F2 = 64
    F3 = 128
    F4 = 256
    F5 = 512
    
    # Input layer
    X_input = Input(input_shape, dtype=tf.float32, name="input")
    
    # Reshape input to 4-D
    X = Reshape((input_shape[0],input_shape[1],1))(X_input)
    
    # Encoder L1
    X = conv_block(X, F1, filt_k, do_rate)
    
    X_skip1 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L2
    X = conv_block(X, F2, filt_k, do_rate)
    
    X_skip2 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L3
    X = conv_block(X, F3, filt_k, do_rate)
    
    X_skip3 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L4
    X = conv_block(X, F4, filt_k, do_rate)
    
    X_skip4 = X
    
    X = MaxPooling2D(pool_size=(2,2), strides=2, padding='same')(X)
    
    # Encoder L5
    X = conv_block(X, F5, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F4, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L4
    X = Concatenate(axis=3)([X,X_skip4])
   
    X = conv_block(X, F4, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F3, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L3
    X = Concatenate(axis=3)([X,X_skip3])
   
    X = conv_block(X, F3, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F2, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L2
    X = Concatenate(axis=3)([X,X_skip2])
   
    X = conv_block(X, F2, filt_k, do_rate)
    
    X = Conv2DTranspose(filters=F1, kernel_size=(filt_k,filt_k), strides=(2,2), padding='same')(X)
    X = Activation('relu')(X)
    
    # Decoder L1    
    X = Concatenate(axis=3)([X,X_skip1])
   
    X = conv_block(X, F1, filt_k, do_rate)
    
    # Output layer
    X = Conv2D(filters=1, kernel_size=(1,1), strides=1, padding='same')(X)
    
    # Reshape output
    X = Reshape((input_shape[0],input_shape[1]))(X)
    
    # Create model
    model = Model(inputs = X_input, outputs = X, name='UNet5_postproc')
    
    return model