text: # coding=utf-8
  # Copyright 2021 The Deeplab2 Authors.
  #
  # 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.
  
  """This file contains functions to build encoder and decoder."""
  import tensorflow as tf
  
  from deeplab2 import config_pb2
  from deeplab2.model.decoder import deeplabv3
  from deeplab2.model.decoder import deeplabv3plus
  from deeplab2.model.decoder import max_deeplab
  from deeplab2.model.decoder import motion_deeplab_decoder
  from deeplab2.model.decoder import panoptic_deeplab
  from deeplab2.model.decoder import vip_deeplab_decoder
  from deeplab2.model.encoder import axial_resnet_instances
  from deeplab2.model.encoder import mobilenet
  
  
  def create_encoder(backbone_options: config_pb2.ModelOptions.BackboneOptions,
                     bn_layer: tf.keras.layers.Layer,
                     conv_kernel_weight_decay: float = 0.0) -> tf.keras.Model:
    """Creates an encoder.
  
    Args:
      backbone_options: A proto config of type
        config_pb2.ModelOptions.BackboneOptions.
      bn_layer: A tf.keras.layers.Layer that computes the normalization.
      conv_kernel_weight_decay: A float, the weight decay for convolution kernels.
  
    Returns:
      An instance of tf.keras.Model containing the encoder.
  
    Raises:
      ValueError: An error occurs when the specified encoder meta architecture is
        not supported.
    """
    if ('resnet' in backbone_options.name or
        'swidernet' in backbone_options.name or
        'axial_deeplab' in backbone_options.name or
        'max_deeplab' in backbone_options.name):
      return create_resnet_encoder(
          backbone_options,
          bn_layer=bn_layer,
          conv_kernel_weight_decay=conv_kernel_weight_decay)
    elif 'mobilenet' in backbone_options.name:
      return create_mobilenet_encoder(
          backbone_options,
          bn_layer=bn_layer,
          conv_kernel_weight_decay=conv_kernel_weight_decay)
    raise ValueError('The specified encoder %s is not a valid encoder.' %
                     backbone_options.name)
  
  
  def create_mobilenet_encoder(
      backbone_options: config_pb2.ModelOptions.BackboneOptions,
      bn_layer: tf.keras.layers.Layer,
      conv_kernel_weight_decay: float = 0.0) -> tf.keras.Model:
    """Creates a MobileNet encoder specified by name.
  
    Args:
      backbone_options: A proto config of type
        config_pb2.ModelOptions.BackboneOptions.
      bn_layer: A tf.keras.layers.Layer that computes the normalization.
      conv_kernel_weight_decay: A float, the weight decay for convolution kernels.
  
    Returns:
      An instance of tf.keras.Model containing the MobileNet encoder.
    """
    if backbone_options.name.lower() == 'mobilenet_v3_large':
      backbone = mobilenet.MobileNetV3Large
    elif backbone_options.name.lower() == 'mobilenet_v3_small':
      backbone = mobilenet.MobileNetV3Small
    else:
      raise ValueError('The specified encoder %s is not a valid encoder.' %
                       backbone_options.name)
    assert backbone_options.use_squeeze_and_excite
    assert backbone_options.drop_path_keep_prob == 1
    assert backbone_options.use_sac_beyond_stride == -1
    assert backbone_options.backbone_layer_multiplier == 1
    return backbone(
        output_stride=backbone_options.output_stride,
        width_multiplier=backbone_options.backbone_width_multiplier,
        bn_layer=bn_layer,
        conv_kernel_weight_decay=conv_kernel_weight_decay)
  
  
  def create_resnet_encoder(
      backbone_options: config_pb2.ModelOptions.BackboneOptions,
      bn_layer: tf.keras.layers.Layer,
      conv_kernel_weight_decay: float = 0.0) -> tf.keras.Model:
    """Creates a ResNet encoder specified by name.
  
    Args:
      backbone_options: A proto config of type
        config_pb2.ModelOptions.BackboneOptions.
      bn_layer: A tf.keras.layers.Layer that computes the normalization.
      conv_kernel_weight_decay: A float, the weight decay for convolution kernels.
  
    Returns:
      An instance of tf.keras.Model containing the ResNet encoder.
    """
    return axial_resnet_instances.get_model(
        backbone_options.name,
        output_stride=backbone_options.output_stride,
        stem_width_multiplier=backbone_options.stem_width_multiplier,
        width_multiplier=backbone_options.backbone_width_multiplier,
        backbone_layer_multiplier=backbone_options.backbone_layer_multiplier,
        block_group_config={
            'use_squeeze_and_excite': backbone_options.use_squeeze_and_excite,
            'drop_path_keep_prob': backbone_options.drop_path_keep_prob,
            'drop_path_schedule': backbone_options.drop_path_schedule,
            'use_sac_beyond_stride': backbone_options.use_sac_beyond_stride},
        bn_layer=bn_layer,
        conv_kernel_weight_decay=conv_kernel_weight_decay)
  
  
  def create_decoder(model_options: config_pb2.ModelOptions,
                     bn_layer: tf.keras.layers.Layer,
                     ignore_label: int) -> tf.keras.Model:
    """Creates a DeepLab decoder.
  
    Args:
      model_options: A proto config of type config_pb2.ModelOptions.
      bn_layer: A tf.keras.layers.Layer that computes the normalization.
      ignore_label: An integer specifying the ignore label.
  
    Returns:
      An instance of tf.keras.layers.Layer containing the decoder.
  
    Raises:
      ValueError: An error occurs when the specified meta architecture is not
        supported.
    """
    meta_architecture = model_options.WhichOneof('meta_architecture')
    if meta_architecture == 'deeplab_v3':
      return deeplabv3.DeepLabV3(
          model_options.decoder, model_options.deeplab_v3, bn_layer=bn_layer)
    elif meta_architecture == 'deeplab_v3_plus':
      return deeplabv3plus.DeepLabV3Plus(
          model_options.decoder, model_options.deeplab_v3_plus, bn_layer=bn_layer)
    elif meta_architecture == 'panoptic_deeplab':
      return panoptic_deeplab.PanopticDeepLab(
          model_options.decoder,
          model_options.panoptic_deeplab,
          bn_layer=bn_layer)
    elif meta_architecture == 'motion_deeplab':
      return motion_deeplab_decoder.MotionDeepLabDecoder(
          model_options.decoder,
          model_options.motion_deeplab,
          bn_layer=bn_layer)
    elif meta_architecture == 'vip_deeplab':
      return vip_deeplab_decoder.ViPDeepLabDecoder(
          model_options.decoder,
          model_options.vip_deeplab,
          bn_layer=bn_layer)
    elif meta_architecture == 'max_deeplab':
      return max_deeplab.MaXDeepLab(
          model_options.decoder,
          model_options.max_deeplab,
          ignore_label=ignore_label,
          bn_layer=bn_layer)
    raise ValueError('The specified meta architecture %s is not implemented.' %
                     meta_architecture)
  