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 code to build a DeepLabV3.17 18Reference:19  - [Rethinking Atrous Convolution for Semantic Image Segmentation](20      https://arxiv.org/pdf/1706.05587.pdf)21"""22import tensorflow as tf23 24from deeplab2 import common25from deeplab2.model.decoder import aspp26from deeplab2.model.layers import convolutions27 28 29layers = tf.keras.layers30 31 32class DeepLabV3(layers.Layer):33  """A DeepLabV3 model.34 35  This model takes in features from an encoder and performs multi-scale context36  aggregation with the help of an ASPP layer. Finally, a classification head is37  used to predict a semantic segmentation.38  """39 40  def __init__(self,41               decoder_options,42               deeplabv3_options,43               bn_layer=tf.keras.layers.BatchNormalization):44    """Creates a DeepLabV3 decoder of type layers.Layer.45 46    Args:47      decoder_options: Decoder options as defined in config_pb2.DecoderOptions.48      deeplabv3_options: Model options as defined in49        config_pb2.ModelOptions.DeeplabV3Options.50      bn_layer: An optional tf.keras.layers.Layer that computes the51        normalization (default: tf.keras.layers.BatchNormalization).52    """53    super(DeepLabV3, self).__init__(name='DeepLabV3')54 55    self._feature_name = decoder_options.feature_key56    self._aspp = aspp.ASPP(decoder_options.aspp_channels,57                           decoder_options.atrous_rates,58                           bn_layer=bn_layer)59 60    self._classifier_conv_bn_act = convolutions.Conv2DSame(61        decoder_options.decoder_channels,62        kernel_size=3,63        name='classifier_conv_bn_act',64        use_bias=False,65        use_bn=True,66        bn_layer=bn_layer,67        activation='relu')68 69    self._final_conv = convolutions.Conv2DSame(70        deeplabv3_options.num_classes, kernel_size=1, name='final_conv')71 72  def set_pool_size(self, pool_size):73    """Sets the pooling size of the ASPP pooling layer.74 75    Args:76      pool_size: A tuple specifying the pooling size of the ASPP pooling layer.77    """78    self._aspp.set_pool_size(pool_size)79 80  def get_pool_size(self):81    return self._aspp.get_pool_size()82 83  def reset_pooling_layer(self):84    """Resets the ASPP pooling layer to global average pooling."""85    self._aspp.reset_pooling_layer()86 87  def call(self, features, training=False):88    """Performs a forward pass.89 90    Args:91      features: A single input tf.Tensor or an input dict of tf.Tensor with92        shape [batch, height, width, channels]. If passed a dict, different keys93        should point to different features extracted by the encoder, e.g.94        low-level or high-level features.95      training: A boolean flag indicating whether training behavior should be96        used (default: False).97 98    Returns:99      A dictionary containing the semantic prediction under key100      common.PRED_SEMANTIC_LOGITS_KEY.101    """102    if isinstance(features, tf.Tensor):103      feature = features104    else:105      feature = features[self._feature_name]106 107    x = self._aspp(feature, training=training)108 109    x = self._classifier_conv_bn_act(x, training=training)110 111    return {common.PRED_SEMANTIC_LOGITS_KEY: self._final_conv(x)}112 113  @property114  def checkpoint_items(self):115    items = {116        common.CKPT_DEEPLABV3_ASPP: self._aspp,117        common.CKPT_DEEPLABV3_CLASSIFIER_CONV_BN_ACT:118            self._classifier_conv_bn_act,119        common.CKPT_SEMANTIC_LAST_LAYER: self._final_conv,120    }121    return items122