|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| r"""Provides DeepLab model definition and helper functions.
|
|
|
| DeepLab is a deep learning system for semantic image segmentation with
|
| the following features:
|
|
|
| (1) Atrous convolution to explicitly control the resolution at which
|
| feature responses are computed within Deep Convolutional Neural Networks.
|
|
|
| (2) Atrous spatial pyramid pooling (ASPP) to robustly segment objects at
|
| multiple scales with filters at multiple sampling rates and effective
|
| fields-of-views.
|
|
|
| (3) ASPP module augmented with image-level feature and batch normalization.
|
|
|
| (4) A simple yet effective decoder module to recover the object boundaries.
|
|
|
| See the following papers for more details:
|
|
|
| "Encoder-Decoder with Atrous Separable Convolution for Semantic Image
|
| Segmentation"
|
| Liang-Chieh Chen, Yukun Zhu, George Papandreou, Florian Schroff, Hartwig Adam.
|
| (https://arxiv.org/abs/1802.02611)
|
|
|
| "Rethinking Atrous Convolution for Semantic Image Segmentation,"
|
| Liang-Chieh Chen, George Papandreou, Florian Schroff, Hartwig Adam
|
| (https://arxiv.org/abs/1706.05587)
|
|
|
| "DeepLab: Semantic Image Segmentation with Deep Convolutional Nets,
|
| Atrous Convolution, and Fully Connected CRFs",
|
| Liang-Chieh Chen*, George Papandreou*, Iasonas Kokkinos, Kevin Murphy,
|
| Alan L Yuille (* equal contribution)
|
| (https://arxiv.org/abs/1606.00915)
|
|
|
| "Semantic Image Segmentation with Deep Convolutional Nets and Fully Connected
|
| CRFs"
|
| Liang-Chieh Chen*, George Papandreou*, Iasonas Kokkinos, Kevin Murphy,
|
| Alan L. Yuille (* equal contribution)
|
| (https://arxiv.org/abs/1412.7062)
|
| """
|
| import tensorflow as tf
|
| from tensorflow.contrib import slim as contrib_slim
|
| from deeplab.core import dense_prediction_cell
|
| from deeplab.core import feature_extractor
|
| from deeplab.core import utils
|
|
|
| slim = contrib_slim
|
|
|
| LOGITS_SCOPE_NAME = 'logits'
|
| MERGED_LOGITS_SCOPE = 'merged_logits'
|
| IMAGE_POOLING_SCOPE = 'image_pooling'
|
| ASPP_SCOPE = 'aspp'
|
| CONCAT_PROJECTION_SCOPE = 'concat_projection'
|
| DECODER_SCOPE = 'decoder'
|
| META_ARCHITECTURE_SCOPE = 'meta_architecture'
|
|
|
| PROB_SUFFIX = '_prob'
|
|
|
| _resize_bilinear = utils.resize_bilinear
|
| scale_dimension = utils.scale_dimension
|
| split_separable_conv2d = utils.split_separable_conv2d
|
|
|
|
|
| def get_extra_layer_scopes(last_layers_contain_logits_only=False):
|
| """Gets the scopes for extra layers.
|
|
|
| Args:
|
| last_layers_contain_logits_only: Boolean, True if only consider logits as
|
| the last layer (i.e., exclude ASPP module, decoder module and so on)
|
|
|
| Returns:
|
| A list of scopes for extra layers.
|
| """
|
| if last_layers_contain_logits_only:
|
| return [LOGITS_SCOPE_NAME]
|
| else:
|
| return [
|
| LOGITS_SCOPE_NAME,
|
| IMAGE_POOLING_SCOPE,
|
| ASPP_SCOPE,
|
| CONCAT_PROJECTION_SCOPE,
|
| DECODER_SCOPE,
|
| META_ARCHITECTURE_SCOPE,
|
| ]
|
|
|
|
|
| def predict_labels_multi_scale(images,
|
| model_options,
|
| eval_scales=(1.0,),
|
| add_flipped_images=False):
|
| """Predicts segmentation labels.
|
|
|
| Args:
|
| images: A tensor of size [batch, height, width, channels].
|
| model_options: A ModelOptions instance to configure models.
|
| eval_scales: The scales to resize images for evaluation.
|
| add_flipped_images: Add flipped images for evaluation or not.
|
|
|
| Returns:
|
| A dictionary with keys specifying the output_type (e.g., semantic
|
| prediction) and values storing Tensors representing predictions (argmax
|
| over channels). Each prediction has size [batch, height, width].
|
| """
|
| outputs_to_predictions = {
|
| output: []
|
| for output in model_options.outputs_to_num_classes
|
| }
|
|
|
| for i, image_scale in enumerate(eval_scales):
|
| with tf.variable_scope(tf.get_variable_scope(), reuse=True if i else None):
|
| outputs_to_scales_to_logits = multi_scale_logits(
|
| images,
|
| model_options=model_options,
|
| image_pyramid=[image_scale],
|
| is_training=False,
|
| fine_tune_batch_norm=False)
|
|
|
| if add_flipped_images:
|
| with tf.variable_scope(tf.get_variable_scope(), reuse=True):
|
| outputs_to_scales_to_logits_reversed = multi_scale_logits(
|
| tf.reverse_v2(images, [2]),
|
| model_options=model_options,
|
| image_pyramid=[image_scale],
|
| is_training=False,
|
| fine_tune_batch_norm=False)
|
|
|
| for output in sorted(outputs_to_scales_to_logits):
|
| scales_to_logits = outputs_to_scales_to_logits[output]
|
| logits = _resize_bilinear(
|
| scales_to_logits[MERGED_LOGITS_SCOPE],
|
| tf.shape(images)[1:3],
|
| scales_to_logits[MERGED_LOGITS_SCOPE].dtype)
|
| outputs_to_predictions[output].append(
|
| tf.expand_dims(tf.nn.softmax(logits), 4))
|
|
|
| if add_flipped_images:
|
| scales_to_logits_reversed = (
|
| outputs_to_scales_to_logits_reversed[output])
|
| logits_reversed = _resize_bilinear(
|
| tf.reverse_v2(scales_to_logits_reversed[MERGED_LOGITS_SCOPE], [2]),
|
| tf.shape(images)[1:3],
|
| scales_to_logits_reversed[MERGED_LOGITS_SCOPE].dtype)
|
| outputs_to_predictions[output].append(
|
| tf.expand_dims(tf.nn.softmax(logits_reversed), 4))
|
|
|
| for output in sorted(outputs_to_predictions):
|
| predictions = outputs_to_predictions[output]
|
|
|
| predictions = tf.reduce_mean(tf.concat(predictions, 4), axis=4)
|
| outputs_to_predictions[output] = tf.argmax(predictions, 3)
|
| outputs_to_predictions[output + PROB_SUFFIX] = tf.nn.softmax(predictions)
|
|
|
| return outputs_to_predictions
|
|
|
|
|
| def predict_labels(images, model_options, image_pyramid=None):
|
| """Predicts segmentation labels.
|
|
|
| Args:
|
| images: A tensor of size [batch, height, width, channels].
|
| model_options: A ModelOptions instance to configure models.
|
| image_pyramid: Input image scales for multi-scale feature extraction.
|
|
|
| Returns:
|
| A dictionary with keys specifying the output_type (e.g., semantic
|
| prediction) and values storing Tensors representing predictions (argmax
|
| over channels). Each prediction has size [batch, height, width].
|
| """
|
| outputs_to_scales_to_logits = multi_scale_logits(
|
| images,
|
| model_options=model_options,
|
| image_pyramid=image_pyramid,
|
| is_training=False,
|
| fine_tune_batch_norm=False)
|
|
|
| predictions = {}
|
| for output in sorted(outputs_to_scales_to_logits):
|
| scales_to_logits = outputs_to_scales_to_logits[output]
|
| logits = scales_to_logits[MERGED_LOGITS_SCOPE]
|
|
|
|
|
|
|
|
|
| if model_options.prediction_with_upsampled_logits:
|
| logits = _resize_bilinear(logits,
|
| tf.shape(images)[1:3],
|
| scales_to_logits[MERGED_LOGITS_SCOPE].dtype)
|
| predictions[output] = tf.argmax(logits, 3)
|
| predictions[output + PROB_SUFFIX] = tf.nn.softmax(logits)
|
| else:
|
| argmax_results = tf.argmax(logits, 3)
|
| argmax_results = tf.image.resize_nearest_neighbor(
|
| tf.expand_dims(argmax_results, 3),
|
| tf.shape(images)[1:3],
|
| align_corners=True,
|
| name='resize_prediction')
|
| predictions[output] = tf.squeeze(argmax_results, 3)
|
| predictions[output + PROB_SUFFIX] = tf.image.resize_bilinear(
|
| tf.nn.softmax(logits),
|
| tf.shape(images)[1:3],
|
| align_corners=True,
|
| name='resize_prob')
|
| return predictions
|
|
|
|
|
| def multi_scale_logits(images,
|
| model_options,
|
| image_pyramid,
|
| weight_decay=0.0001,
|
| is_training=False,
|
| fine_tune_batch_norm=False,
|
| nas_training_hyper_parameters=None):
|
| """Gets the logits for multi-scale inputs.
|
|
|
| The returned logits are all downsampled (due to max-pooling layers)
|
| for both training and evaluation.
|
|
|
| Args:
|
| images: A tensor of size [batch, height, width, channels].
|
| model_options: A ModelOptions instance to configure models.
|
| image_pyramid: Input image scales for multi-scale feature extraction.
|
| weight_decay: The weight decay for model variables.
|
| is_training: Is training or not.
|
| fine_tune_batch_norm: Fine-tune the batch norm parameters or not.
|
| nas_training_hyper_parameters: A dictionary storing hyper-parameters for
|
| training nas models. Its keys are:
|
| - `drop_path_keep_prob`: Probability to keep each path in the cell when
|
| training.
|
| - `total_training_steps`: Total training steps to help drop path
|
| probability calculation.
|
|
|
| Returns:
|
| outputs_to_scales_to_logits: A map of maps from output_type (e.g.,
|
| semantic prediction) to a dictionary of multi-scale logits names to
|
| logits. For each output_type, the dictionary has keys which
|
| correspond to the scales and values which correspond to the logits.
|
| For example, if `scales` equals [1.0, 1.5], then the keys would
|
| include 'merged_logits', 'logits_1.00' and 'logits_1.50'.
|
|
|
| Raises:
|
| ValueError: If model_options doesn't specify crop_size and its
|
| add_image_level_feature = True, since add_image_level_feature requires
|
| crop_size information.
|
| """
|
|
|
| if not image_pyramid:
|
| image_pyramid = [1.0]
|
| crop_height = (
|
| model_options.crop_size[0]
|
| if model_options.crop_size else tf.shape(images)[1])
|
| crop_width = (
|
| model_options.crop_size[1]
|
| if model_options.crop_size else tf.shape(images)[2])
|
| if model_options.image_pooling_crop_size:
|
| image_pooling_crop_height = model_options.image_pooling_crop_size[0]
|
| image_pooling_crop_width = model_options.image_pooling_crop_size[1]
|
|
|
|
|
| if model_options.decoder_output_stride:
|
| logits_output_stride = min(model_options.decoder_output_stride)
|
| else:
|
| logits_output_stride = model_options.output_stride
|
|
|
| logits_height = scale_dimension(
|
| crop_height,
|
| max(1.0, max(image_pyramid)) / logits_output_stride)
|
| logits_width = scale_dimension(
|
| crop_width,
|
| max(1.0, max(image_pyramid)) / logits_output_stride)
|
|
|
|
|
| outputs_to_scales_to_logits = {
|
| k: {}
|
| for k in model_options.outputs_to_num_classes
|
| }
|
|
|
| num_channels = images.get_shape().as_list()[-1]
|
|
|
| for image_scale in image_pyramid:
|
| if image_scale != 1.0:
|
| scaled_height = scale_dimension(crop_height, image_scale)
|
| scaled_width = scale_dimension(crop_width, image_scale)
|
| scaled_crop_size = [scaled_height, scaled_width]
|
| scaled_images = _resize_bilinear(images, scaled_crop_size, images.dtype)
|
| if model_options.crop_size:
|
| scaled_images.set_shape(
|
| [None, scaled_height, scaled_width, num_channels])
|
|
|
| scaled_image_pooling_crop_size = None
|
| if model_options.image_pooling_crop_size:
|
| scaled_image_pooling_crop_size = [
|
| scale_dimension(image_pooling_crop_height, image_scale),
|
| scale_dimension(image_pooling_crop_width, image_scale)]
|
| else:
|
| scaled_crop_size = model_options.crop_size
|
| scaled_images = images
|
| scaled_image_pooling_crop_size = model_options.image_pooling_crop_size
|
|
|
| updated_options = model_options._replace(
|
| crop_size=scaled_crop_size,
|
| image_pooling_crop_size=scaled_image_pooling_crop_size)
|
| outputs_to_logits = _get_logits(
|
| scaled_images,
|
| updated_options,
|
| weight_decay=weight_decay,
|
| reuse=tf.AUTO_REUSE,
|
| is_training=is_training,
|
| fine_tune_batch_norm=fine_tune_batch_norm,
|
| nas_training_hyper_parameters=nas_training_hyper_parameters)
|
|
|
|
|
| for output in sorted(outputs_to_logits):
|
| outputs_to_logits[output] = _resize_bilinear(
|
| outputs_to_logits[output], [logits_height, logits_width],
|
| outputs_to_logits[output].dtype)
|
|
|
|
|
| if len(image_pyramid) == 1:
|
| for output in sorted(model_options.outputs_to_num_classes):
|
| outputs_to_scales_to_logits[output][
|
| MERGED_LOGITS_SCOPE] = outputs_to_logits[output]
|
| return outputs_to_scales_to_logits
|
|
|
|
|
| for output in sorted(model_options.outputs_to_num_classes):
|
| outputs_to_scales_to_logits[output][
|
| 'logits_%.2f' % image_scale] = outputs_to_logits[output]
|
|
|
|
|
| for output in sorted(model_options.outputs_to_num_classes):
|
|
|
| all_logits = [
|
| tf.expand_dims(logits, axis=4)
|
| for logits in outputs_to_scales_to_logits[output].values()
|
| ]
|
| all_logits = tf.concat(all_logits, 4)
|
| merge_fn = (
|
| tf.reduce_max
|
| if model_options.merge_method == 'max' else tf.reduce_mean)
|
| outputs_to_scales_to_logits[output][MERGED_LOGITS_SCOPE] = merge_fn(
|
| all_logits, axis=4)
|
|
|
| return outputs_to_scales_to_logits
|
|
|
|
|
| def extract_features(images,
|
| model_options,
|
| weight_decay=0.0001,
|
| reuse=None,
|
| is_training=False,
|
| fine_tune_batch_norm=False,
|
| nas_training_hyper_parameters=None):
|
| """Extracts features by the particular model_variant.
|
|
|
| Args:
|
| images: A tensor of size [batch, height, width, channels].
|
| model_options: A ModelOptions instance to configure models.
|
| weight_decay: The weight decay for model variables.
|
| reuse: Reuse the model variables or not.
|
| is_training: Is training or not.
|
| fine_tune_batch_norm: Fine-tune the batch norm parameters or not.
|
| nas_training_hyper_parameters: A dictionary storing hyper-parameters for
|
| training nas models. Its keys are:
|
| - `drop_path_keep_prob`: Probability to keep each path in the cell when
|
| training.
|
| - `total_training_steps`: Total training steps to help drop path
|
| probability calculation.
|
|
|
| Returns:
|
| concat_logits: A tensor of size [batch, feature_height, feature_width,
|
| feature_channels], where feature_height/feature_width are determined by
|
| the images height/width and output_stride.
|
| end_points: A dictionary from components of the network to the corresponding
|
| activation.
|
| """
|
| features, end_points = feature_extractor.extract_features(
|
| images,
|
| output_stride=model_options.output_stride,
|
| multi_grid=model_options.multi_grid,
|
| model_variant=model_options.model_variant,
|
| depth_multiplier=model_options.depth_multiplier,
|
| divisible_by=model_options.divisible_by,
|
| weight_decay=weight_decay,
|
| reuse=reuse,
|
| is_training=is_training,
|
| preprocessed_images_dtype=model_options.preprocessed_images_dtype,
|
| fine_tune_batch_norm=fine_tune_batch_norm,
|
| nas_architecture_options=model_options.nas_architecture_options,
|
| nas_training_hyper_parameters=nas_training_hyper_parameters,
|
| use_bounded_activation=model_options.use_bounded_activation)
|
|
|
| if not model_options.aspp_with_batch_norm:
|
| return features, end_points
|
| else:
|
| if model_options.dense_prediction_cell_config is not None:
|
| tf.logging.info('Using dense prediction cell config.')
|
| dense_prediction_layer = dense_prediction_cell.DensePredictionCell(
|
| config=model_options.dense_prediction_cell_config,
|
| hparams={
|
| 'conv_rate_multiplier': 16 // model_options.output_stride,
|
| })
|
| concat_logits = dense_prediction_layer.build_cell(
|
| features,
|
| output_stride=model_options.output_stride,
|
| crop_size=model_options.crop_size,
|
| image_pooling_crop_size=model_options.image_pooling_crop_size,
|
| weight_decay=weight_decay,
|
| reuse=reuse,
|
| is_training=is_training,
|
| fine_tune_batch_norm=fine_tune_batch_norm)
|
| return concat_logits, end_points
|
| else:
|
|
|
|
|
|
|
|
|
| batch_norm_params = utils.get_batch_norm_params(
|
| decay=0.9997,
|
| epsilon=1e-5,
|
| scale=True,
|
| is_training=(is_training and fine_tune_batch_norm),
|
| sync_batch_norm_method=model_options.sync_batch_norm_method)
|
| batch_norm = utils.get_batch_norm_fn(
|
| model_options.sync_batch_norm_method)
|
| activation_fn = (
|
| tf.nn.relu6 if model_options.use_bounded_activation else tf.nn.relu)
|
| with slim.arg_scope(
|
| [slim.conv2d, slim.separable_conv2d],
|
| weights_regularizer=slim.l2_regularizer(weight_decay),
|
| activation_fn=activation_fn,
|
| normalizer_fn=batch_norm,
|
| padding='SAME',
|
| stride=1,
|
| reuse=reuse):
|
| with slim.arg_scope([batch_norm], **batch_norm_params):
|
| depth = model_options.aspp_convs_filters
|
| branch_logits = []
|
|
|
| if model_options.add_image_level_feature:
|
| if model_options.crop_size is not None:
|
| image_pooling_crop_size = model_options.image_pooling_crop_size
|
|
|
| if image_pooling_crop_size is None:
|
| image_pooling_crop_size = model_options.crop_size
|
| pool_height = scale_dimension(
|
| image_pooling_crop_size[0],
|
| 1. / model_options.output_stride)
|
| pool_width = scale_dimension(
|
| image_pooling_crop_size[1],
|
| 1. / model_options.output_stride)
|
| image_feature = slim.avg_pool2d(
|
| features, [pool_height, pool_width],
|
| model_options.image_pooling_stride, padding='VALID')
|
| resize_height = scale_dimension(
|
| model_options.crop_size[0],
|
| 1. / model_options.output_stride)
|
| resize_width = scale_dimension(
|
| model_options.crop_size[1],
|
| 1. / model_options.output_stride)
|
| else:
|
|
|
| pool_height = tf.shape(features)[1]
|
| pool_width = tf.shape(features)[2]
|
| image_feature = tf.reduce_mean(
|
| features, axis=[1, 2], keepdims=True)
|
| resize_height = pool_height
|
| resize_width = pool_width
|
| image_feature_activation_fn = tf.nn.relu
|
| image_feature_normalizer_fn = batch_norm
|
| if model_options.aspp_with_squeeze_and_excitation:
|
| image_feature_activation_fn = tf.nn.sigmoid
|
| if model_options.image_se_uses_qsigmoid:
|
| image_feature_activation_fn = utils.q_sigmoid
|
| image_feature_normalizer_fn = None
|
| image_feature = slim.conv2d(
|
| image_feature, depth, 1,
|
| activation_fn=image_feature_activation_fn,
|
| normalizer_fn=image_feature_normalizer_fn,
|
| scope=IMAGE_POOLING_SCOPE)
|
| image_feature = _resize_bilinear(
|
| image_feature,
|
| [resize_height, resize_width],
|
| image_feature.dtype)
|
|
|
| if isinstance(resize_height, tf.Tensor):
|
| resize_height = None
|
| if isinstance(resize_width, tf.Tensor):
|
| resize_width = None
|
| image_feature.set_shape([None, resize_height, resize_width, depth])
|
| if not model_options.aspp_with_squeeze_and_excitation:
|
| branch_logits.append(image_feature)
|
|
|
|
|
| branch_logits.append(slim.conv2d(features, depth, 1,
|
| scope=ASPP_SCOPE + str(0)))
|
|
|
| if model_options.atrous_rates:
|
|
|
| for i, rate in enumerate(model_options.atrous_rates, 1):
|
| scope = ASPP_SCOPE + str(i)
|
| if model_options.aspp_with_separable_conv:
|
| aspp_features = split_separable_conv2d(
|
| features,
|
| filters=depth,
|
| rate=rate,
|
| weight_decay=weight_decay,
|
| scope=scope)
|
| else:
|
| aspp_features = slim.conv2d(
|
| features, depth, 3, rate=rate, scope=scope)
|
| branch_logits.append(aspp_features)
|
|
|
|
|
| concat_logits = tf.concat(branch_logits, 3)
|
| if model_options.aspp_with_concat_projection:
|
| concat_logits = slim.conv2d(
|
| concat_logits, depth, 1, scope=CONCAT_PROJECTION_SCOPE)
|
| concat_logits = slim.dropout(
|
| concat_logits,
|
| keep_prob=0.9,
|
| is_training=is_training,
|
| scope=CONCAT_PROJECTION_SCOPE + '_dropout')
|
| if (model_options.add_image_level_feature and
|
| model_options.aspp_with_squeeze_and_excitation):
|
| concat_logits *= image_feature
|
|
|
| return concat_logits, end_points
|
|
|
|
|
| def _get_logits(images,
|
| model_options,
|
| weight_decay=0.0001,
|
| reuse=None,
|
| is_training=False,
|
| fine_tune_batch_norm=False,
|
| nas_training_hyper_parameters=None):
|
| """Gets the logits by atrous/image spatial pyramid pooling.
|
|
|
| Args:
|
| images: A tensor of size [batch, height, width, channels].
|
| model_options: A ModelOptions instance to configure models.
|
| weight_decay: The weight decay for model variables.
|
| reuse: Reuse the model variables or not.
|
| is_training: Is training or not.
|
| fine_tune_batch_norm: Fine-tune the batch norm parameters or not.
|
| nas_training_hyper_parameters: A dictionary storing hyper-parameters for
|
| training nas models. Its keys are:
|
| - `drop_path_keep_prob`: Probability to keep each path in the cell when
|
| training.
|
| - `total_training_steps`: Total training steps to help drop path
|
| probability calculation.
|
|
|
| Returns:
|
| outputs_to_logits: A map from output_type to logits.
|
| """
|
| features, end_points = extract_features(
|
| images,
|
| model_options,
|
| weight_decay=weight_decay,
|
| reuse=reuse,
|
| is_training=is_training,
|
| fine_tune_batch_norm=fine_tune_batch_norm,
|
| nas_training_hyper_parameters=nas_training_hyper_parameters)
|
|
|
| if model_options.decoder_output_stride:
|
| crop_size = model_options.crop_size
|
| if crop_size is None:
|
| crop_size = [tf.shape(images)[1], tf.shape(images)[2]]
|
| features = refine_by_decoder(
|
| features,
|
| end_points,
|
| crop_size=crop_size,
|
| decoder_output_stride=model_options.decoder_output_stride,
|
| decoder_use_separable_conv=model_options.decoder_use_separable_conv,
|
| decoder_use_sum_merge=model_options.decoder_use_sum_merge,
|
| decoder_filters=model_options.decoder_filters,
|
| decoder_output_is_logits=model_options.decoder_output_is_logits,
|
| model_variant=model_options.model_variant,
|
| weight_decay=weight_decay,
|
| reuse=reuse,
|
| is_training=is_training,
|
| fine_tune_batch_norm=fine_tune_batch_norm,
|
| use_bounded_activation=model_options.use_bounded_activation)
|
|
|
| outputs_to_logits = {}
|
| for output in sorted(model_options.outputs_to_num_classes):
|
| if model_options.decoder_output_is_logits:
|
| outputs_to_logits[output] = tf.identity(features,
|
| name=output)
|
| else:
|
| outputs_to_logits[output] = get_branch_logits(
|
| features,
|
| model_options.outputs_to_num_classes[output],
|
| model_options.atrous_rates,
|
| aspp_with_batch_norm=model_options.aspp_with_batch_norm,
|
| kernel_size=model_options.logits_kernel_size,
|
| weight_decay=weight_decay,
|
| reuse=reuse,
|
| scope_suffix=output)
|
|
|
| return outputs_to_logits
|
|
|
|
|
| def refine_by_decoder(features,
|
| end_points,
|
| crop_size=None,
|
| decoder_output_stride=None,
|
| decoder_use_separable_conv=False,
|
| decoder_use_sum_merge=False,
|
| decoder_filters=256,
|
| decoder_output_is_logits=False,
|
| model_variant=None,
|
| weight_decay=0.0001,
|
| reuse=None,
|
| is_training=False,
|
| fine_tune_batch_norm=False,
|
| use_bounded_activation=False,
|
| sync_batch_norm_method='None'):
|
| """Adds the decoder to obtain sharper segmentation results.
|
|
|
| Args:
|
| features: A tensor of size [batch, features_height, features_width,
|
| features_channels].
|
| end_points: A dictionary from components of the network to the corresponding
|
| activation.
|
| crop_size: A tuple [crop_height, crop_width] specifying whole patch crop
|
| size.
|
| decoder_output_stride: A list of integers specifying the output stride of
|
| low-level features used in the decoder module.
|
| decoder_use_separable_conv: Employ separable convolution for decoder or not.
|
| decoder_use_sum_merge: Boolean, decoder uses simple sum merge or not.
|
| decoder_filters: Integer, decoder filter size.
|
| decoder_output_is_logits: Boolean, using decoder output as logits or not.
|
| model_variant: Model variant for feature extraction.
|
| weight_decay: The weight decay for model variables.
|
| reuse: Reuse the model variables or not.
|
| is_training: Is training or not.
|
| fine_tune_batch_norm: Fine-tune the batch norm parameters or not.
|
| use_bounded_activation: Whether or not to use bounded activations. Bounded
|
| activations better lend themselves to quantized inference.
|
| sync_batch_norm_method: String, method used to sync batch norm. Currently
|
| only support `None` (no sync batch norm) and `tpu` (use tpu code to
|
| sync batch norm).
|
|
|
| Returns:
|
| Decoder output with size [batch, decoder_height, decoder_width,
|
| decoder_channels].
|
|
|
| Raises:
|
| ValueError: If crop_size is None.
|
| """
|
| if crop_size is None:
|
| raise ValueError('crop_size must be provided when using decoder.')
|
| batch_norm_params = utils.get_batch_norm_params(
|
| decay=0.9997,
|
| epsilon=1e-5,
|
| scale=True,
|
| is_training=(is_training and fine_tune_batch_norm),
|
| sync_batch_norm_method=sync_batch_norm_method)
|
| batch_norm = utils.get_batch_norm_fn(sync_batch_norm_method)
|
| decoder_depth = decoder_filters
|
| projected_filters = 48
|
| if decoder_use_sum_merge:
|
|
|
|
|
| projected_filters = decoder_filters
|
| if decoder_output_is_logits:
|
|
|
| activation_fn = None
|
| normalizer_fn = None
|
| conv2d_kernel = 1
|
|
|
| decoder_use_separable_conv = False
|
| else:
|
|
|
| activation_fn = tf.nn.relu6 if use_bounded_activation else tf.nn.relu
|
| normalizer_fn = batch_norm
|
| conv2d_kernel = 3
|
| with slim.arg_scope(
|
| [slim.conv2d, slim.separable_conv2d],
|
| weights_regularizer=slim.l2_regularizer(weight_decay),
|
| activation_fn=activation_fn,
|
| normalizer_fn=normalizer_fn,
|
| padding='SAME',
|
| stride=1,
|
| reuse=reuse):
|
| with slim.arg_scope([batch_norm], **batch_norm_params):
|
| with tf.variable_scope(DECODER_SCOPE, DECODER_SCOPE, [features]):
|
| decoder_features = features
|
| decoder_stage = 0
|
| scope_suffix = ''
|
| for output_stride in decoder_output_stride:
|
| feature_list = feature_extractor.networks_to_feature_maps[
|
| model_variant][
|
| feature_extractor.DECODER_END_POINTS][output_stride]
|
|
|
|
|
| if decoder_stage:
|
| scope_suffix = '_{}'.format(decoder_stage)
|
| for i, name in enumerate(feature_list):
|
| decoder_features_list = [decoder_features]
|
|
|
| if ('mobilenet' in model_variant or
|
| model_variant.startswith('mnas') or
|
| model_variant.startswith('nas')):
|
| feature_name = name
|
| else:
|
| feature_name = '{}/{}'.format(
|
| feature_extractor.name_scope[model_variant], name)
|
| decoder_features_list.append(
|
| slim.conv2d(
|
| end_points[feature_name],
|
| projected_filters,
|
| 1,
|
| scope='feature_projection' + str(i) + scope_suffix))
|
|
|
| decoder_height = scale_dimension(crop_size[0], 1.0 / output_stride)
|
| decoder_width = scale_dimension(crop_size[1], 1.0 / output_stride)
|
|
|
| for j, feature in enumerate(decoder_features_list):
|
| decoder_features_list[j] = _resize_bilinear(
|
| feature, [decoder_height, decoder_width], feature.dtype)
|
| h = (None if isinstance(decoder_height, tf.Tensor)
|
| else decoder_height)
|
| w = (None if isinstance(decoder_width, tf.Tensor)
|
| else decoder_width)
|
| decoder_features_list[j].set_shape([None, h, w, None])
|
| if decoder_use_sum_merge:
|
| decoder_features = _decoder_with_sum_merge(
|
| decoder_features_list,
|
| decoder_depth,
|
| conv2d_kernel=conv2d_kernel,
|
| decoder_use_separable_conv=decoder_use_separable_conv,
|
| weight_decay=weight_decay,
|
| scope_suffix=scope_suffix)
|
| else:
|
| if not decoder_use_separable_conv:
|
| scope_suffix = str(i) + scope_suffix
|
| decoder_features = _decoder_with_concat_merge(
|
| decoder_features_list,
|
| decoder_depth,
|
| decoder_use_separable_conv=decoder_use_separable_conv,
|
| weight_decay=weight_decay,
|
| scope_suffix=scope_suffix)
|
| decoder_stage += 1
|
| return decoder_features
|
|
|
|
|
| def _decoder_with_sum_merge(decoder_features_list,
|
| decoder_depth,
|
| conv2d_kernel=3,
|
| decoder_use_separable_conv=True,
|
| weight_decay=0.0001,
|
| scope_suffix=''):
|
| """Decoder with sum to merge features.
|
|
|
| Args:
|
| decoder_features_list: A list of decoder features.
|
| decoder_depth: Integer, the filters used in the convolution.
|
| conv2d_kernel: Integer, the convolution kernel size.
|
| decoder_use_separable_conv: Boolean, use separable conv or not.
|
| weight_decay: Weight decay for the model variables.
|
| scope_suffix: String, used in the scope suffix.
|
|
|
| Returns:
|
| decoder features merged with sum.
|
|
|
| Raises:
|
| RuntimeError: If decoder_features_list have length not equal to 2.
|
| """
|
| if len(decoder_features_list) != 2:
|
| raise RuntimeError('Expect decoder_features has length 2.')
|
|
|
| if decoder_use_separable_conv:
|
| decoder_features = split_separable_conv2d(
|
| decoder_features_list[0],
|
| filters=decoder_depth,
|
| rate=1,
|
| weight_decay=weight_decay,
|
| scope='decoder_split_sep_conv0'+scope_suffix) + decoder_features_list[1]
|
| else:
|
| decoder_features = slim.conv2d(
|
| decoder_features_list[0],
|
| decoder_depth,
|
| conv2d_kernel,
|
| scope='decoder_conv0'+scope_suffix) + decoder_features_list[1]
|
| return decoder_features
|
|
|
|
|
| def _decoder_with_concat_merge(decoder_features_list,
|
| decoder_depth,
|
| decoder_use_separable_conv=True,
|
| weight_decay=0.0001,
|
| scope_suffix=''):
|
| """Decoder with concatenation to merge features.
|
|
|
| This decoder method applies two convolutions to smooth the features obtained
|
| by concatenating the input decoder_features_list.
|
|
|
| This decoder module is proposed in the DeepLabv3+ paper.
|
|
|
| Args:
|
| decoder_features_list: A list of decoder features.
|
| decoder_depth: Integer, the filters used in the convolution.
|
| decoder_use_separable_conv: Boolean, use separable conv or not.
|
| weight_decay: Weight decay for the model variables.
|
| scope_suffix: String, used in the scope suffix.
|
|
|
| Returns:
|
| decoder features merged with concatenation.
|
| """
|
| if decoder_use_separable_conv:
|
| decoder_features = split_separable_conv2d(
|
| tf.concat(decoder_features_list, 3),
|
| filters=decoder_depth,
|
| rate=1,
|
| weight_decay=weight_decay,
|
| scope='decoder_conv0'+scope_suffix)
|
| decoder_features = split_separable_conv2d(
|
| decoder_features,
|
| filters=decoder_depth,
|
| rate=1,
|
| weight_decay=weight_decay,
|
| scope='decoder_conv1'+scope_suffix)
|
| else:
|
| num_convs = 2
|
| decoder_features = slim.repeat(
|
| tf.concat(decoder_features_list, 3),
|
| num_convs,
|
| slim.conv2d,
|
| decoder_depth,
|
| 3,
|
| scope='decoder_conv'+scope_suffix)
|
| return decoder_features
|
|
|
|
|
| def get_branch_logits(features,
|
| num_classes,
|
| atrous_rates=None,
|
| aspp_with_batch_norm=False,
|
| kernel_size=1,
|
| weight_decay=0.0001,
|
| reuse=None,
|
| scope_suffix=''):
|
| """Gets the logits from each model's branch.
|
|
|
| The underlying model is branched out in the last layer when atrous
|
| spatial pyramid pooling is employed, and all branches are sum-merged
|
| to form the final logits.
|
|
|
| Args:
|
| features: A float tensor of shape [batch, height, width, channels].
|
| num_classes: Number of classes to predict.
|
| atrous_rates: A list of atrous convolution rates for last layer.
|
| aspp_with_batch_norm: Use batch normalization layers for ASPP.
|
| kernel_size: Kernel size for convolution.
|
| weight_decay: Weight decay for the model variables.
|
| reuse: Reuse model variables or not.
|
| scope_suffix: Scope suffix for the model variables.
|
|
|
| Returns:
|
| Merged logits with shape [batch, height, width, num_classes].
|
|
|
| Raises:
|
| ValueError: Upon invalid input kernel_size value.
|
| """
|
|
|
|
|
| if aspp_with_batch_norm or atrous_rates is None:
|
| if kernel_size != 1:
|
| raise ValueError('Kernel size must be 1 when atrous_rates is None or '
|
| 'using aspp_with_batch_norm. Gets %d.' % kernel_size)
|
| atrous_rates = [1]
|
|
|
| with slim.arg_scope(
|
| [slim.conv2d],
|
| weights_regularizer=slim.l2_regularizer(weight_decay),
|
| weights_initializer=tf.truncated_normal_initializer(stddev=0.01),
|
| reuse=reuse):
|
| with tf.variable_scope(LOGITS_SCOPE_NAME, LOGITS_SCOPE_NAME, [features]):
|
| branch_logits = []
|
| for i, rate in enumerate(atrous_rates):
|
| scope = scope_suffix
|
| if i:
|
| scope += '_%d' % i
|
|
|
| branch_logits.append(
|
| slim.conv2d(
|
| features,
|
| num_classes,
|
| kernel_size=kernel_size,
|
| rate=rate,
|
| activation_fn=None,
|
| normalizer_fn=None,
|
| scope=scope))
|
|
|
| return tf.add_n(branch_logits)
|
|
|
