karolmajek/maxdeeplab

1# coding=utf-82# Copyright 2021 The Deeplab2 Authors.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16"""This file contains functions to build encoder and decoder."""17import tensorflow as tf18 19from deeplab2 import config_pb220from deeplab2.model.decoder import deeplabv321from deeplab2.model.decoder import deeplabv3plus22from deeplab2.model.decoder import max_deeplab23from deeplab2.model.decoder import motion_deeplab_decoder24from deeplab2.model.decoder import panoptic_deeplab25from deeplab2.model.decoder import vip_deeplab_decoder26from deeplab2.model.encoder import axial_resnet_instances27from deeplab2.model.encoder import mobilenet28 29 30def create_encoder(backbone_options: config_pb2.ModelOptions.BackboneOptions,31                   bn_layer: tf.keras.layers.Layer,32                   conv_kernel_weight_decay: float = 0.0) -> tf.keras.Model:33  """Creates an encoder.34 35  Args:36    backbone_options: A proto config of type37      config_pb2.ModelOptions.BackboneOptions.38    bn_layer: A tf.keras.layers.Layer that computes the normalization.39    conv_kernel_weight_decay: A float, the weight decay for convolution kernels.40 41  Returns:42    An instance of tf.keras.Model containing the encoder.43 44  Raises:45    ValueError: An error occurs when the specified encoder meta architecture is46      not supported.47  """48  if ('resnet' in backbone_options.name or49      'swidernet' in backbone_options.name or50      'axial_deeplab' in backbone_options.name or51      'max_deeplab' in backbone_options.name):52    return create_resnet_encoder(53        backbone_options,54        bn_layer=bn_layer,55        conv_kernel_weight_decay=conv_kernel_weight_decay)56  elif 'mobilenet' in backbone_options.name:57    return create_mobilenet_encoder(58        backbone_options,59        bn_layer=bn_layer,60        conv_kernel_weight_decay=conv_kernel_weight_decay)61  raise ValueError('The specified encoder %s is not a valid encoder.' %62                   backbone_options.name)63 64 65def create_mobilenet_encoder(66    backbone_options: config_pb2.ModelOptions.BackboneOptions,67    bn_layer: tf.keras.layers.Layer,68    conv_kernel_weight_decay: float = 0.0) -> tf.keras.Model:69  """Creates a MobileNet encoder specified by name.70 71  Args:72    backbone_options: A proto config of type73      config_pb2.ModelOptions.BackboneOptions.74    bn_layer: A tf.keras.layers.Layer that computes the normalization.75    conv_kernel_weight_decay: A float, the weight decay for convolution kernels.76 77  Returns:78    An instance of tf.keras.Model containing the MobileNet encoder.79  """80  if backbone_options.name.lower() == 'mobilenet_v3_large':81    backbone = mobilenet.MobileNetV3Large82  elif backbone_options.name.lower() == 'mobilenet_v3_small':83    backbone = mobilenet.MobileNetV3Small84  else:85    raise ValueError('The specified encoder %s is not a valid encoder.' %86                     backbone_options.name)87  assert backbone_options.use_squeeze_and_excite88  assert backbone_options.drop_path_keep_prob == 189  assert backbone_options.use_sac_beyond_stride == -190  assert backbone_options.backbone_layer_multiplier == 191  return backbone(92      output_stride=backbone_options.output_stride,93      width_multiplier=backbone_options.backbone_width_multiplier,94      bn_layer=bn_layer,95      conv_kernel_weight_decay=conv_kernel_weight_decay)96 97 98def create_resnet_encoder(99    backbone_options: config_pb2.ModelOptions.BackboneOptions,100    bn_layer: tf.keras.layers.Layer,101    conv_kernel_weight_decay: float = 0.0) -> tf.keras.Model:102  """Creates a ResNet encoder specified by name.103 104  Args:105    backbone_options: A proto config of type106      config_pb2.ModelOptions.BackboneOptions.107    bn_layer: A tf.keras.layers.Layer that computes the normalization.108    conv_kernel_weight_decay: A float, the weight decay for convolution kernels.109 110  Returns:111    An instance of tf.keras.Model containing the ResNet encoder.112  """113  return axial_resnet_instances.get_model(114      backbone_options.name,115      output_stride=backbone_options.output_stride,116      stem_width_multiplier=backbone_options.stem_width_multiplier,117      width_multiplier=backbone_options.backbone_width_multiplier,118      backbone_layer_multiplier=backbone_options.backbone_layer_multiplier,119      block_group_config={120          'use_squeeze_and_excite': backbone_options.use_squeeze_and_excite,121          'drop_path_keep_prob': backbone_options.drop_path_keep_prob,122          'drop_path_schedule': backbone_options.drop_path_schedule,123          'use_sac_beyond_stride': backbone_options.use_sac_beyond_stride},124      bn_layer=bn_layer,125      conv_kernel_weight_decay=conv_kernel_weight_decay)126 127 128def create_decoder(model_options: config_pb2.ModelOptions,129                   bn_layer: tf.keras.layers.Layer,130                   ignore_label: int) -> tf.keras.Model:131  """Creates a DeepLab decoder.132 133  Args:134    model_options: A proto config of type config_pb2.ModelOptions.135    bn_layer: A tf.keras.layers.Layer that computes the normalization.136    ignore_label: An integer specifying the ignore label.137 138  Returns:139    An instance of tf.keras.layers.Layer containing the decoder.140 141  Raises:142    ValueError: An error occurs when the specified meta architecture is not143      supported.144  """145  meta_architecture = model_options.WhichOneof('meta_architecture')146  if meta_architecture == 'deeplab_v3':147    return deeplabv3.DeepLabV3(148        model_options.decoder, model_options.deeplab_v3, bn_layer=bn_layer)149  elif meta_architecture == 'deeplab_v3_plus':150    return deeplabv3plus.DeepLabV3Plus(151        model_options.decoder, model_options.deeplab_v3_plus, bn_layer=bn_layer)152  elif meta_architecture == 'panoptic_deeplab':153    return panoptic_deeplab.PanopticDeepLab(154        model_options.decoder,155        model_options.panoptic_deeplab,156        bn_layer=bn_layer)157  elif meta_architecture == 'motion_deeplab':158    return motion_deeplab_decoder.MotionDeepLabDecoder(159        model_options.decoder,160        model_options.motion_deeplab,161        bn_layer=bn_layer)162  elif meta_architecture == 'vip_deeplab':163    return vip_deeplab_decoder.ViPDeepLabDecoder(164        model_options.decoder,165        model_options.vip_deeplab,166        bn_layer=bn_layer)167  elif meta_architecture == 'max_deeplab':168    return max_deeplab.MaXDeepLab(169        model_options.decoder,170        model_options.max_deeplab,171        ignore_label=ignore_label,172        bn_layer=bn_layer)173  raise ValueError('The specified meta architecture %s is not implemented.' %174                   meta_architecture)175