#' 2-D implementation of the U-net deep learning architecture.
#'
#' Creates a keras model of the U-net deep learning architecture for image
#' segmentation and regression. More information is provided at the authors'
#' website:
#'
#' \url{https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/}
#'
#' with the paper available here:
#'
#' \url{https://arxiv.org/abs/1505.04597}
#'
#' This particular implementation was influenced by the following python
#' implementation:
#'
#' \url{https://github.com/joelthelion/ultrasound-nerve-segmentation}
#'
#' @param inputImageSize Used for specifying the input tensor shape. The
#' shape (or dimension) of that tensor is the image dimensions followed by
#' the number of channels (e.g., red, green, and blue). The batch size
#' (i.e., number of training images) is not specified a priori.
#' @param numberOfOutputs Meaning depends on the \code{mode}. For
#' 'classification' this is the number of segmentation labels. For 'regression'
#' this is the number of outputs.
#' @param scalarOutputSize if greater than 0, a global average pooling from each
#' encoding layer is concatenated to a dense layer as a secondary output.
#' @param scalarOutputActivation activation for nonzero output scalar.
#' @param numberOfLayers number of encoding/decoding layers.
#' @param numberOfFiltersAtBaseLayer number of filters at the beginning and end
#' of the \verb{'U'}. Doubles at each descending/ascending layer.
#' @param numberOfFilters vector explicitly setting the number of filters at
#' each layer. One can either set this or \code{numberOfLayers} and
#' \code{numberOfFiltersAtBaseLayer}. Default = NULL.
#' @param convolutionKernelSize 2-d vector defining the kernel size
#' during the encoding path.
#' @param deconvolutionKernelSize 2-d vector defining the kernel size
#' during the decoding.
#' @param poolSize 2-d vector defining the region for each pooling layer.
#' @param strides 2-d vector describing the stride length in each direction.
#' @param dropoutRate float between 0 and 1 to use between dense layers.
#' @param weightDecay weighting parameter for L2 regularization of the
#' kernel weights of the convolution layers. Default = 0.0.
#' @param additionalOptions string or vector of strings specifying
#' specific configuration add-ons/tweaks:
#' \itemize{
#' \item{"attentionGating"}{attention-unet variant in
#' \url{https://pubmed.ncbi.nlm.nih.gov/33288961/}}
#' \item{"nnUnetActivationStyle"}{U-net activation explained in
#' \url{https://pubmed.ncbi.nlm.nih.gov/33288961/}}
#' \item{"initialConvolutionKernelSize[X]"}{Set the first two convolutional layer
#' kernel sizes to X.}
#' }
#' @param mode 'classification' or 'regression' or 'sigmoid'.
#'
#' @return a u-net keras model
#' @author Tustison NJ
#' @examples
#'
#' library( ANTsR )
#' library( ANTsRNet )
#' library( keras )
#'
#' imageIDs <- c( "r16", "r27", "r30", "r62", "r64", "r85" )
#' trainingBatchSize <- length( imageIDs )
#'
#' # Perform simple 3-tissue segmentation.
#'
#' segmentationLabels <- c( 1, 2, 3 )
#' numberOfLabels <- length( segmentationLabels )
#' initialization <- paste0( 'KMeans[', numberOfLabels, ']' )
#'
#' domainImage <- antsImageRead( getANTsRData( imageIDs[1] ) )
#'
#' X_train <- array( data = NA, dim = c( trainingBatchSize, dim( domainImage ), 1 ) )
#' Y_train <- array( data = NA, dim = c( trainingBatchSize, dim( domainImage ) ) )
#'
#' images <- list()
#' segmentations <- list()
#'
#' for( i in seq_len( trainingBatchSize ) )
#' {
#' cat( "Processing image", imageIDs[i], "\n" )
#' image <- antsImageRead( getANTsRData( imageIDs[i] ) )
#' mask <- getMask( image )
#' segmentation <- atropos( image, mask, initialization )$segmentation
#'
#' X_train[i,,, 1] <- as.array( image )
#' Y_train[i,,] <- as.array( segmentation )
#' }
#' Y_train <- encodeUnet( Y_train, segmentationLabels )
#'
#' # Perform a simple normalization
#'
#' X_train <- ( X_train - mean( X_train ) ) / sd( X_train )
#'
#' # Create the model
#'
#' model <- createUnetModel2D( c( dim( domainImage ), 1 ),
#' numberOfOutputs = numberOfLabels )
#'
#' metric_multilabel_dice_coefficient <-
#' custom_metric( "multilabel_dice_coefficient",
#' multilabel_dice_coefficient )
#'
#' loss_dice <- function( y_true, y_pred ) {
#' -multilabel_dice_coefficient(y_true, y_pred)
#' }
#' attr(loss_dice, "py_function_name") <- "multilabel_dice_coefficient"
#'
#' model %>% compile( loss = loss_dice,
#' optimizer = optimizer_adam( lr = 0.0001 ),
#' metrics = metric_multilabel_dice_coefficient )
#'
#' # Comment out the rest due to travis build constraints
#'
#' # Fit the model
#'
#' # track <- model %>% fit( X_train, Y_train,
#' # epochs = 100, batch_size = 5, verbose = 1, shuffle = TRUE,
#' # callbacks = list(
#' # callback_model_checkpoint( "unetModelInterimWeights.h5",
#' # monitor = 'val_loss', save_best_only = TRUE ),
#' # callback_reduce_lr_on_plateau( monitor = "val_loss", factor = 0.1 )
#' # ),
#' # validation_split = 0.2 )
#'
#' # Save the model and/or save the model weights
#'
#' # save_model_hdf5( model, filepath = 'unetModel.h5' )
#' # save_model_weights_hdf5( unetModel, filepath = 'unetModelWeights.h5' ) )
#'
#' @import keras
#' @export
createUnetModel2D <- function( inputImageSize,
numberOfOutputs = 2,
scalarOutputSize = 0,
scalarOutputActivation = "relu",
numberOfLayers = 4,
numberOfFiltersAtBaseLayer = 32,
numberOfFilters = NULL,
convolutionKernelSize = c( 3, 3 ),
deconvolutionKernelSize = c( 2, 2 ),
poolSize = c( 2, 2 ),
strides = c( 2, 2 ),
dropoutRate = 0.0,
weightDecay = 0.0,
mode = c( 'classification', 'regression', 'sigmoid' ),
additionalOptions = NA
)
{
nnUnetActivation <- function( x )
{
x <- x %>% layer_instance_normalization()
x <- x %>% layer_activation_leaky_relu( alpha = 0.01 )
return( x )
}
attentionGate2D <- function( x, g, interShape )
{
xTheta <- x %>% layer_conv_2d( filters = interShape, kernel_size = c( 1L, 1L ),
strides = c( 1L, 1L ) )
gPhi <- g %>% layer_conv_2d( filters = interShape, kernel_size = c( 1L, 1L ),
strides = c( 1L, 1L ) )
f <- layer_add( list( xTheta, gPhi ) )
f <- f %>% layer_activation_relu()
fPsi <- f %>% layer_conv_2d( filters = 1L, kernel_size = c( 1L, 1L ),
strides = c( 1L, 1L ) )
alpha <- fPsi %>% layer_activation( activation = "sigmoid" )
attention <- layer_multiply( list( x, alpha ) )
return( attention )
}
mode <- match.arg( mode )
# Handle additional specific configurations
initialConvolutionKernelSize <- convolutionKernelSize
addAttentionGating <- FALSE
nnUnetActivationStyle <- FALSE
if( ! any( is.na( additionalOptions ) ) )
{
if( "attentionGating" %in% additionalOptions )
{
addAttentionGating <- TRUE
}
if( "nnUnetActivationStyle" %in% additionalOptions )
{
nnUnetActivationStyle <- TRUE
}
index <- pmatch( "initialConvolutionKernelSize", additionalOptions )
if( ! is.na( index ) )
{
option <- additionalOptions[index]
initialConvolutionKernelSize <- strtoi( sub( "\\]", "", strsplit( option, "\\[" )[[1]][2] ) )
}
}
# Specify the number of filters
if( ! is.null( numberOfFilters ) )
{
numberOfLayers <- length( numberOfFilters )
} else {
numberOfFilters <- rep( 0, numberOfLayers )
for( i in seq_len( numberOfLayers ) )
{
numberOfFilters[i] <- numberOfFiltersAtBaseLayer * 2 ^ ( i - 1 )
}
}
inputs <- layer_input( shape = inputImageSize )
# Encoding path
encodingConvolutionLayers <- list()
for( i in seq_len( numberOfLayers ) )
{
if( i == 1 )
{
conv <- inputs %>% layer_conv_2d( filters = numberOfFilters[i],
kernel_size = initialConvolutionKernelSize, padding = 'same',
kernel_regularizer = regularizer_l2( weightDecay ) )
} else {
conv <- pool %>% layer_conv_2d( filters = numberOfFilters[i],
kernel_size = convolutionKernelSize, padding = 'same',
kernel_regularizer = regularizer_l2( weightDecay ) )
}
if( nnUnetActivationStyle )
{
conv <- nnUnetActivation( conv )
} else {
conv <- conv %>% layer_activation_relu()
}
if( dropoutRate > 0.0 )
{
conv <- conv %>% layer_dropout( rate = dropoutRate )
}
if( i == 1 )
{
conv <- conv %>% layer_conv_2d( filters = numberOfFilters[i],
kernel_size = initialConvolutionKernelSize, padding = 'same' )
} else {
conv <- conv %>% layer_conv_2d( filters = numberOfFilters[i],
kernel_size = convolutionKernelSize, padding = 'same' )
}
if( nnUnetActivationStyle )
{
conv <- nnUnetActivation( conv )
} else {
conv <- conv %>% layer_activation_relu()
}
encodingConvolutionLayers[[i]] <- conv
if( i < numberOfLayers )
{
pool <- encodingConvolutionLayers[[i]] %>%
layer_max_pooling_2d( pool_size = poolSize, strides = strides,
padding = 'same' )
}
}
scalarOutput <- NULL
if( scalarOutputSize > 0 )
{
scalarLayers <- list()
for( i in seq.int( numberOfLayers ) )
{
scalarLayers[[i]] <- encodingConvolutionLayers[[i]] %>%
layer_global_average_pooling_2d()
}
scalarOutput <- layer_concatenate( scalarLayers, trainable = TRUE ) %>%
layer_dense( units = scalarOutputSize, activation = scalarOutputActivation )
}
# Decoding path
outputs <- encodingConvolutionLayers[[numberOfLayers]]
for( i in 2:numberOfLayers )
{
deconv <- outputs %>%
layer_conv_2d_transpose( filters = numberOfFilters[numberOfLayers - i + 1],
kernel_size = deconvolutionKernelSize,
padding = 'same',
kernel_regularizer = regularizer_l2( weightDecay ) )
if( nnUnetActivationStyle )
{
deconv <- nnUnetActivation( deconv )
}
deconv <- deconv %>% layer_upsampling_2d( size = poolSize )
if( addAttentionGating == TRUE )
{
outputs <- attentionGate2D( deconv,
encodingConvolutionLayers[[numberOfLayers - i + 1]],
as.integer( numberOfFilters[numberOfLayers - i + 1] / 4 ) )
outputs <- layer_concatenate( list( deconv, outputs ), axis = 3, trainable = TRUE )
} else {
outputs <- layer_concatenate( list( deconv,
encodingConvolutionLayers[[numberOfLayers - i + 1]] ),
axis = 3, trainable = TRUE )
}
outputs <- outputs %>%
layer_conv_2d( filters = numberOfFilters[numberOfLayers - i + 1],
kernel_size = convolutionKernelSize, padding = 'same',
kernel_regularizer = regularizer_l2( weightDecay ) )
if( nnUnetActivationStyle )
{
outputs <- nnUnetActivation( outputs )
} else {
outputs <- outputs %>% layer_activation_relu()
}
if( dropoutRate > 0.0 )
{
outputs <- outputs %>% layer_dropout( rate = dropoutRate )
}
outputs <- outputs %>%
layer_conv_2d( filters = numberOfFilters[numberOfLayers - i + 1],
kernel_size = convolutionKernelSize, padding = 'same',
kernel_regularizer = regularizer_l2( weightDecay ) )
if( nnUnetActivationStyle )
{
outputs <- nnUnetActivation( outputs )
} else {
outputs <- outputs %>% layer_activation_relu()
}
}
convActivation <- ''
if( mode == 'sigmoid' )
{
convActivation <- 'sigmoid'
} else if( mode == 'classification' ) {
convActivation <- 'softmax'
} else if( mode == 'regression' ) {
convActivation <- 'linear'
} else {
stop( 'Error: unrecognized mode.' )
}
outputs <- outputs %>%
layer_conv_2d( filters = numberOfOutputs,
kernel_size = c( 1, 1 ), activation = convActivation,
kernel_regularizer = regularizer_l2( weightDecay ) )
unetModel <- NULL
if( scalarOutputSize > 0 )
{
unetModel <- keras_model( inputs = inputs, outputs = list( outputs, scalarOutput ) )
} else {
unetModel <- keras_model( inputs = inputs, outputs = outputs )
}
return( unetModel )
}
#' 3-D image segmentation implementation of the U-net deep learning architecture.
#'
#' Creates a keras model of the U-net deep learning architecture for image
#' segmentation and regression. More information is provided at the authors'
#' website:
#'
#' \url{https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/}
#'
#' with the paper available here:
#'
#' \url{https://arxiv.org/abs/1505.04597}
#'
#' This particular implementation was influenced by the following python
#' implementation:
#'
#' \url{https://github.com/joelthelion/ultrasound-nerve-segmentation}
#'
#' @param inputImageSize Used for specifying the input tensor shape. The
#' shape (or dimension) of that tensor is the image dimensions followed by
#' the number of channels (e.g., red, green, and blue). The batch size
#' (i.e., number of training images) is not specified a priori.
#' @param numberOfOutputs Meaning depends on the \code{mode}. For
#' 'classification' this is the number of segmentation labels. For 'regression'
#' this is the number of outputs.
#' @param scalarOutputSize if greater than 0, a global average pooling from each
#' encoding layer is concatenated to a dense layer as a secondary output.
#' @param scalarOutputActivation activation for nonzero output scalar.
#' @param numberOfLayers number of encoding/decoding layers.
#' @param numberOfFiltersAtBaseLayer number of filters at the beginning and end
#' of the \verb{'U'}. Doubles at each descending/ascending layer.
#' @param numberOfFilters vector explicitly setting the number of filters at
#' each layer. One can either set this or \code{numberOfLayers} and
#' @param convolutionKernelSize 3-d vector defining the kernel size
#' during the encoding path.
#' @param deconvolutionKernelSize 3-d vector defining the kernel size
#' during the decoding.
#' @param poolSize 3-d vector defining the region for each pooling layer.
#' @param strides 3-d vector describing the stride length in each direction.
#' @param dropoutRate float between 0 and 1 to use between dense layers.
#' @param weightDecay weighting parameter for L2 regularization of the
#' kernel weights of the convolution layers. Default = 0.0.
#' @param mode 'classification' or 'regression' or 'sigmoid'.
#' @param additionalOptions string or vector of strings specifying
#' specific configuration add-ons/tweaks:
#' \itemize{
#' \item{"attentionGating"}{attention-unet variant in
#' \url{https://pubmed.ncbi.nlm.nih.gov/33288961/}}
#' \item{"nnUnetActivationStyle"}{U-net activation explained in
#' \url{https://pubmed.ncbi.nlm.nih.gov/33288961/}}
#' \item{"initialConvolutionalKernelSize[X]"}{Set the first two convolutional layer
#' kernel sizes to X.}
#' }
#'
#' @return a u-net keras model
#' @author Tustison NJ
#' @examples
#' # Simple examples, must run successfully and quickly. These will be tested.
#'
#' library( ANTsRNet )
#' library( keras )
#'
#' model <- createUnetModel3D( c( 64, 64, 64, 1 ) )
#'
#' metric_multilabel_dice_coefficient <-
#' custom_metric( "multilabel_dice_coefficient",
#' multilabel_dice_coefficient )
#'
#' loss_dice <- function( y_true, y_pred ) {
#' -multilabel_dice_coefficient(y_true, y_pred)
#' }
#' attr(loss_dice, "py_function_name") <- "multilabel_dice_coefficient"
#'
#' model %>% compile( loss = loss_dice,
#' optimizer = optimizer_adam( lr = 0.0001 ),
#' metrics = metric_multilabel_dice_coefficient )
#'
#' print( model )
#'
#' @import keras
#' @export
createUnetModel3D <- function( inputImageSize,
numberOfOutputs = 2,
scalarOutputSize = 0,
scalarOutputActivation = "relu",
numberOfLayers = 4,
numberOfFiltersAtBaseLayer = 32,
numberOfFilters = NULL,
convolutionKernelSize = c( 3, 3, 3 ),
deconvolutionKernelSize = c( 2, 2, 2 ),
poolSize = c( 2, 2, 2 ),
strides = c( 2, 2, 2 ),
dropoutRate = 0.0,
weightDecay = 0.0,
mode = c( 'classification', 'regression', 'sigmoid' ),
additionalOptions = NA
)
{
nnUnetActivation <- function( x )
{
x <- x %>% layer_instance_normalization()
x <- x %>% layer_activation_leaky_relu( alpha = 0.01 )
return( x )
}
attentionGate3D <- function( x, g, interShape )
{
xTheta <- x %>% layer_conv_3d( filters = interShape, kernel_size = c( 1L, 1L, 1L ),
strides = c( 1L, 1L, 1L ) )
gPhi <- g %>% layer_conv_3d( filters = interShape, kernel_size = c( 1L, 1L, 1L ),
strides = c( 1L, 1L, 1L ) )
f <- layer_add( list( xTheta, gPhi ) )
f <- f %>% layer_activation_relu()
fPsi <- f %>% layer_conv_3d( filters = 1L, kernel_size = c( 1L, 1L, 1L ),
strides = c( 1L, 1L, 1L ) )
alpha <- fPsi %>% layer_activation( activation = "sigmoid" )
attention <- layer_multiply( list( x, alpha ) )
return( attention )
}
mode <- match.arg( mode )
# Handle additional specific configurations
initialConvolutionKernelSize <- convolutionKernelSize
addAttentionGating <- FALSE
nnUnetActivationStyle <- FALSE
if( ! any( is.na( additionalOptions ) ) )
{
if( "attentionGating" %in% additionalOptions )
{
addAttentionGating <- TRUE
}
if( "nnUnetActivationStyle" %in% additionalOptions )
{
nnUnetActivationStyle <- TRUE
}
index <- pmatch( "initialConvolutionKernelSize", additionalOptions )
if( ! is.na( index ) )
{
option <- additionalOptions[index]
initialConvolutionKernelSize <- strtoi( sub( "\\]", "", strsplit( option, "\\[" )[[1]][2] ) )
}
}
# Specify the number of filters
if( ! is.null( numberOfFilters ) )
{
numberOfLayers <- length( numberOfFilters )
} else {
numberOfFilters <- rep( 0, numberOfLayers )
for( i in seq_len( numberOfLayers ) )
{
numberOfFilters[i] <- numberOfFiltersAtBaseLayer * 2 ^ ( i - 1 )
}
}
inputs <- layer_input( shape = inputImageSize )
# Encoding path
encodingConvolutionLayers <- list()
for( i in seq_len( numberOfLayers ) )
{
if( i == 1 )
{
conv <- inputs %>% layer_conv_3d( filters = numberOfFilters[i],
kernel_size = initialConvolutionKernelSize, padding = 'same',
kernel_regularizer = regularizer_l2( weightDecay ) )
} else {
conv <- pool %>% layer_conv_3d( filters = numberOfFilters[i],
kernel_size = convolutionKernelSize, padding = 'same',
kernel_regularizer = regularizer_l2( weightDecay ) )
}
if( nnUnetActivationStyle )
{
conv <- nnUnetActivation( conv )
} else {
conv <- conv %>% layer_activation_relu()
}
if( dropoutRate > 0.0 )
{
conv <- conv %>% layer_dropout( rate = dropoutRate )
}
if( i == 1 )
{
conv <- conv %>% layer_conv_3d( filters = numberOfFilters[i],
kernel_size = initialConvolutionKernelSize, padding = 'same' )
} else {
conv <- conv %>% layer_conv_3d( filters = numberOfFilters[i],
kernel_size = convolutionKernelSize, padding = 'same' )
}
if( nnUnetActivationStyle )
{
conv <- nnUnetActivation( conv )
} else {
conv <- conv %>% layer_activation_relu()
}
encodingConvolutionLayers[[i]] <- conv
if( i < numberOfLayers )
{
pool <- encodingConvolutionLayers[[i]] %>%
layer_max_pooling_3d( pool_size = poolSize, strides = strides,
padding = 'same' )
}
}
scalarOutput <- NULL
if( scalarOutputSize > 0 )
{
scalarLayers <- list()
for( i in seq.int( numberOfLayers ) )
{
scalarLayers[[i]] <- encodingConvolutionLayers[[i]] %>%
layer_global_average_pooling_3d()
}
scalarOutput <- layer_concatenate( scalarLayers, trainable = TRUE ) %>%
layer_dense( units = scalarSize, activation = scalarOutputActivation )
}
# Decoding path
outputs <- encodingConvolutionLayers[[numberOfLayers]]
for( i in 2:numberOfLayers )
{
deconv <- outputs %>%
layer_conv_3d_transpose( filters = numberOfFilters[numberOfLayers - i + 1],
kernel_size = deconvolutionKernelSize,
padding = 'same',
kernel_regularizer = regularizer_l2( weightDecay ) )
if( nnUnetActivationStyle )
{
deconv <- nnUnetActivation( deconv )
}
deconv <- deconv %>% layer_upsampling_3d( size = poolSize )
if( addAttentionGating == TRUE )
{
outputs <- attentionGate3D( deconv,
encodingConvolutionLayers[[numberOfLayers - i + 1]],
as.integer( numberOfFilters[numberOfLayers - i + 1] / 4 ) )
outputs <- layer_concatenate( list( deconv, outputs ), axis = 4, trainable = TRUE )
} else {
outputs <- layer_concatenate( list( deconv,
encodingConvolutionLayers[[numberOfLayers - i + 1]] ),
axis = 4, trainable = TRUE )
}
outputs <- outputs %>%
layer_conv_3d( filters = numberOfFilters[numberOfLayers - i + 1],
kernel_size = convolutionKernelSize, padding = 'same',
kernel_regularizer = regularizer_l2( weightDecay ) )
if( nnUnetActivationStyle )
{
outputs <- nnUnetActivation( outputs )
} else {
outputs <- outputs %>% layer_activation_relu()
}
if( dropoutRate > 0.0 )
{
outputs <- outputs %>% layer_dropout( rate = dropoutRate )
}
outputs <- outputs %>%
layer_conv_3d( filters = numberOfFilters[numberOfLayers - i + 1],
kernel_size = convolutionKernelSize, padding = 'same',
kernel_regularizer = regularizer_l2( weightDecay ) )
if( nnUnetActivationStyle )
{
outputs <- nnUnetActivation( outputs )
} else {
outputs <- outputs %>% layer_activation_relu()
}
}
convActivation <- ''
if( mode == 'sigmoid' )
{
convActivation <- 'sigmoid'
} else if( mode == 'classification' ) {
convActivation <- 'softmax'
} else if( mode == 'regression' ) {
convActivation <- 'linear'
} else {
stop( 'Error: unrecognized mode.' )
}
outputs <- outputs %>%
layer_conv_3d( filters = numberOfOutputs,
kernel_size = c( 1, 1, 1 ), activation = convActivation,
kernel_regularizer = regularizer_l2( weightDecay ) )
unetModel <- NULL
if( scalarOutputSize > 0 )
{
unetModel <- keras_model( inputs = inputs, outputs = list( outputs, scalarOutput ) )
} else {
unetModel <- keras_model( inputs = inputs, outputs = outputs )
}
return( unetModel )
}
Add the following code to your website.
For more information on customizing the embed code, read Embedding Snippets.