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 the code for the Motion-DeepLab decoder."""17 18import tensorflow as tf19 20from deeplab2 import common21from deeplab2 import config_pb222from deeplab2.model.decoder import panoptic_deeplab23 24 25class MotionDeepLabDecoder(tf.keras.layers.Layer):26  """A Motion-DeepLab decoder layer.27 28  This layer takes low- and high-level features as input and uses a dual-ASPP29  and dual-decoder structure to aggregate features for semantic and instance30  segmentation. On top of the decoders, four heads are used to predict semantic31  segmentation, instance center probabilities, instance center regression, and32  previous frame offset regression per pixel.33  """34 35  def __init__(36      self,37      decoder_options: config_pb2.DecoderOptions,38      motion_deeplab_options: config_pb2.ModelOptions.MotionDeepLabOptions,39      bn_layer=tf.keras.layers.BatchNormalization):40    """Initializes a Motion-DeepLab decoder.41 42    Args:43      decoder_options: Decoder options as defined in config_pb2.DecoderOptions.44      motion_deeplab_options: Model options as defined in45        config_pb2.ModelOptions.MotionDeeplabOptions.46      bn_layer: An optional tf.keras.layers.Layer that computes the47        normalization (default: tf.keras.layers.BatchNormalization).48    """49    super(MotionDeepLabDecoder, self).__init__(name='MotionDeepLabDecoder')50 51    low_level_feature_keys = [52        item.feature_key for item in motion_deeplab_options.low_level53    ]54    low_level_channels_project = [55        item.channels_project for item in motion_deeplab_options.low_level56    ]57 58    self._semantic_decoder = panoptic_deeplab.PanopticDeepLabSingleDecoder(59        decoder_options.feature_key,60        low_level_feature_keys,61        low_level_channels_project,62        decoder_options.aspp_channels,63        decoder_options.decoder_channels,64        decoder_options.atrous_rates,65        name='semantic_decoder',66        bn_layer=bn_layer)67    self._semantic_head = panoptic_deeplab.PanopticDeepLabSingleHead(68        motion_deeplab_options.semantic_head.head_channels,69        motion_deeplab_options.semantic_head.output_channels,70        common.PRED_SEMANTIC_LOGITS_KEY,71        name='semantic_head',72        bn_layer=bn_layer)73 74    self._instance_decoder = None75    self._instance_center_head = None76    self._instance_regression_head = None77    self._motion_regression_head = None78 79    if motion_deeplab_options.instance.low_level_override:80      low_level_options = motion_deeplab_options.instance.low_level_override81    else:82      low_level_options = motion_deeplab_options.low_level83 84    # If instance_decoder is set, use those options; otherwise reuse the85    # architecture as defined for the semantic decoder.86    if motion_deeplab_options.instance.HasField('instance_decoder_override'):87      decoder_options = (motion_deeplab_options.instance88                         .instance_decoder_override)89 90    low_level_feature_keys = [item.feature_key for item in low_level_options]91    low_level_channels_project = [92        item.channels_project for item in low_level_options93    ]94 95    self._instance_decoder = panoptic_deeplab.PanopticDeepLabSingleDecoder(96        decoder_options.feature_key,97        low_level_feature_keys,98        low_level_channels_project,99        decoder_options.aspp_channels,100        decoder_options.decoder_channels,101        decoder_options.atrous_rates,102        name='instance_decoder',103        bn_layer=bn_layer)104    self._instance_center_head = panoptic_deeplab.PanopticDeepLabSingleHead(105        motion_deeplab_options.instance.center_head.head_channels,106        motion_deeplab_options.instance.center_head.output_channels,107        common.PRED_CENTER_HEATMAP_KEY,108        name='instance_center_head',109        bn_layer=bn_layer)110    self._instance_regression_head = panoptic_deeplab.PanopticDeepLabSingleHead(111        motion_deeplab_options.instance.regression_head.head_channels,112        motion_deeplab_options.instance.regression_head.output_channels,113        common.PRED_OFFSET_MAP_KEY,114        name='instance_regression_head',115        bn_layer=bn_layer)116 117    # The motion head regresses every pixel to its center in the previous118    # frame.119    self._motion_regression_head = panoptic_deeplab.PanopticDeepLabSingleHead(120        motion_deeplab_options.motion_head.head_channels,121        motion_deeplab_options.motion_head.output_channels,122        common.PRED_FRAME_OFFSET_MAP_KEY,123        name='motion_regression_head',124        bn_layer=bn_layer)125 126  def reset_pooling_layer(self):127    """Resets the ASPP pooling layers to global average pooling."""128    self._semantic_decoder.reset_pooling_layer()129    if self._instance_decoder is not None:130      self._instance_decoder.reset_pooling_layer()131 132  def set_pool_size(self, pool_size):133    """Sets the pooling size of the ASPP pooling layers.134 135    Args:136      pool_size: A tuple specifying the pooling size of the ASPP pooling layers.137    """138    self._semantic_decoder.set_pool_size(pool_size)139    if self._instance_decoder is not None:140      self._instance_decoder.set_pool_size(pool_size)141 142  def get_pool_size(self):143    return self._semantic_decoder.get_pool_size()144 145  def call(self, features, training=False):146    """Performs a forward pass.147 148    Args:149      features: An input dict of tf.Tensor with shape [batch, height, width,150        channels]. Different keys should point to different features extracted151        by the encoder, e.g. low-level or high-level features.152      training: A boolean flag indicating whether training behavior should be153        used (default: False).154 155    Returns:156      A dictionary containing the results of the semantic segmentation head and157        depending on the configuration also of the instance segmentation head.158    """159 160    semantic_features = self._semantic_decoder(features, training=training)161    results = self._semantic_head(semantic_features, training=training)162 163    if self._instance_decoder is not None:164      instance_features = self._instance_decoder(features, training=training)165      instance_center_predictions = self._instance_center_head(166          instance_features, training=training)167      instance_regression_predictions = self._instance_regression_head(168          instance_features, training=training)169      motion_regression_predictions = self._motion_regression_head(170          instance_features, training=training)171      if results.keys() & motion_regression_predictions.keys():172        raise ValueError('The keys of the semantic branch and the instance '173                         'motion branch overlap. Please use unique keys.')174      results.update(motion_regression_predictions)175 176      if results.keys() & instance_center_predictions.keys():177        raise ValueError('The keys of the semantic branch and the instance '178                         'center branch overlap. Please use unique keys.')179      results.update(instance_center_predictions)180 181      if results.keys() & instance_regression_predictions.keys():182        raise ValueError('The keys of the semantic branch and the instance '183                         'regression branch overlap. Please use unique keys.')184      results.update(instance_regression_predictions)185 186    return results187 188  @property189  def checkpoint_items(self):190    items = {191        common.CKPT_SEMANTIC_DECODER:192            self._semantic_decoder,193        common.CKPT_SEMANTIC_HEAD_WITHOUT_LAST_LAYER:194            self._semantic_head.conv_block,195        common.CKPT_SEMANTIC_LAST_LAYER:196            self._semantic_head.final_conv197    }198    if self._instance_decoder is not None:199      instance_items = {200          common.CKPT_INSTANCE_DECODER:201              self._instance_decoder,202          common.CKPT_INSTANCE_CENTER_HEAD_WITHOUT_LAST_LAYER:203              self._instance_center_head.conv_block,204          common.CKPT_INSTANCE_CENTER_HEAD_LAST_LAYER:205              self._instance_center_head.final_conv,206          common.CKPT_INSTANCE_REGRESSION_HEAD_WITHOUT_LAST_LAYER:207              self._instance_regression_head.conv_block,208          common.CKPT_INSTANCE_REGRESSION_HEAD_LAST_LAYER:209              self._instance_regression_head.final_conv,210          common.CKPT_MOTION_REGRESSION_HEAD_WITHOUT_LAST_LAYER:211              self._motion_regression_head.conv_block,212          common.CKPT_MOTION_REGRESSION_HEAD_LAST_LAYER:213              self._motion_regression_head.final_conv,214      }215      items.update(instance_items)216    return items217