karolmajek/maxdeeplab
0
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 ViP-DeepLab decoder.17 18Reference:19 - [ViP-DeepLab: Learning Visual Perception with Depth-aware Video20 Panoptic Segmentation](https://arxiv.org/abs/2012.05258)21"""22import tensorflow as tf23 24from deeplab2 import common25from deeplab2.model.decoder import panoptic_deeplab26 27 28layers = tf.keras.layers29 30 31class ViPDeepLabDecoder(layers.Layer):32 """A ViP-DeepLab decoder layer.33 34 This layer takes low- and high-level features as input and uses a dual-ASPP35 and dual-decoder structure to aggregate features for semantic and instance36 segmentation. On top of the decoders, three heads are used to predict semantic37 segmentation, instance center probabilities, and instance center regression38 per pixel. It also has a branch to predict the next-frame instance center39 regression. Different from the ViP-DeepLab paper which uses Cascade-ASPP, this40 reimplementation only uses ASPP.41 """42 43 def __init__(self,44 decoder_options,45 vip_deeplab_options,46 bn_layer=tf.keras.layers.BatchNormalization):47 """Initializes a ViP-DeepLab decoder.48 49 Args:50 decoder_options: Decoder options as defined in config_pb2.DecoderOptions.51 vip_deeplab_options: Model options as defined in52 config_pb2.ModelOptions.ViPDeeplabOptions.53 bn_layer: An optional tf.keras.layers.Layer that computes the54 normalization (default: tf.keras.layers.BatchNormalization).55 """56 super(ViPDeepLabDecoder, self).__init__(name='ViPDeepLab')57 58 low_level_feature_keys = [59 item.feature_key for item in vip_deeplab_options.low_level60 ]61 low_level_channels_project = [62 item.channels_project for item in vip_deeplab_options.low_level63 ]64 65 self._semantic_decoder = panoptic_deeplab.PanopticDeepLabSingleDecoder(66 high_level_feature_name=decoder_options.feature_key,67 low_level_feature_names=low_level_feature_keys,68 low_level_channels_project=low_level_channels_project,69 aspp_output_channels=decoder_options.aspp_channels,70 decoder_output_channels=decoder_options.decoder_channels,71 atrous_rates=decoder_options.atrous_rates,72 name='semantic_decoder',73 aspp_use_only_1x1_proj_conv=decoder_options.aspp_use_only_1x1_proj_conv,74 decoder_conv_type=decoder_options.decoder_conv_type,75 bn_layer=bn_layer)76 self._semantic_head = panoptic_deeplab.PanopticDeepLabSingleHead(77 vip_deeplab_options.semantic_head.head_channels,78 vip_deeplab_options.semantic_head.output_channels,79 common.PRED_SEMANTIC_LOGITS_KEY,80 name='semantic_head',81 conv_type=vip_deeplab_options.semantic_head.head_conv_type,82 bn_layer=bn_layer)83 84 self._instance_decoder = None85 self._instance_center_head = None86 self._instance_regression_head = None87 self._next_instance_decoder = None88 self._next_instance_regression_head = None89 90 if vip_deeplab_options.instance.enable:91 if vip_deeplab_options.instance.low_level_override:92 low_level_options = vip_deeplab_options.instance.low_level_override93 else:94 low_level_options = vip_deeplab_options.low_level95 96 # If instance_decoder is set, use those options; otherwise reuse the97 # architecture as defined for the semantic decoder.98 if vip_deeplab_options.instance.HasField(99 'instance_decoder_override'):100 decoder_options = (vip_deeplab_options.instance101 .instance_decoder_override)102 103 low_level_feature_keys = [item.feature_key for item in low_level_options]104 low_level_channels_project = [105 item.channels_project for item in low_level_options106 ]107 108 self._instance_decoder = panoptic_deeplab.PanopticDeepLabSingleDecoder(109 high_level_feature_name=decoder_options.feature_key,110 low_level_feature_names=low_level_feature_keys,111 low_level_channels_project=low_level_channels_project,112 aspp_output_channels=decoder_options.aspp_channels,113 decoder_output_channels=decoder_options.decoder_channels,114 atrous_rates=decoder_options.atrous_rates,115 name='instance_decoder',116 aspp_use_only_1x1_proj_conv=(117 decoder_options.aspp_use_only_1x1_proj_conv),118 decoder_conv_type=decoder_options.decoder_conv_type,119 bn_layer=bn_layer)120 self._instance_center_head = panoptic_deeplab.PanopticDeepLabSingleHead(121 vip_deeplab_options.instance.center_head.head_channels,122 vip_deeplab_options.instance.center_head.output_channels,123 common.PRED_CENTER_HEATMAP_KEY,124 name='instance_center_head',125 conv_type=(126 vip_deeplab_options.instance.center_head.head_conv_type),127 bn_layer=bn_layer)128 self._instance_regression_head = (129 panoptic_deeplab.PanopticDeepLabSingleHead(130 vip_deeplab_options.instance.regression_head.head_channels,131 vip_deeplab_options.instance.regression_head.output_channels,132 common.PRED_OFFSET_MAP_KEY,133 name='instance_regression_head',134 conv_type=(135 vip_deeplab_options.instance.regression_head.head_conv_type),136 bn_layer=bn_layer))137 138 if vip_deeplab_options.instance.HasField('next_regression_head'):139 self._next_instance_decoder = (140 panoptic_deeplab.PanopticDeepLabSingleDecoder(141 high_level_feature_name=decoder_options.feature_key,142 low_level_feature_names=low_level_feature_keys,143 low_level_channels_project=low_level_channels_project,144 aspp_output_channels=decoder_options.aspp_channels,145 decoder_output_channels=decoder_options.decoder_channels,146 atrous_rates=decoder_options.atrous_rates,147 name='next_instance_decoder',148 aspp_use_only_1x1_proj_conv=(149 decoder_options.aspp_use_only_1x1_proj_conv),150 decoder_conv_type=decoder_options.decoder_conv_type,151 bn_layer=bn_layer))152 self._next_instance_regression_head = (153 panoptic_deeplab.PanopticDeepLabSingleHead(154 (vip_deeplab_options.instance.next_regression_head155 .head_channels),156 (vip_deeplab_options.instance.next_regression_head157 .output_channels),158 common.PRED_NEXT_OFFSET_MAP_KEY,159 name='next_instance_regression_head',160 conv_type=(vip_deeplab_options.instance.next_regression_head161 .head_conv_type),162 bn_layer=bn_layer))163 self._next_high_level_feature_name = decoder_options.feature_key164 165 def reset_pooling_layer(self):166 """Resets the ASPP pooling layers to global average pooling."""167 self._semantic_decoder.reset_pooling_layer()168 if self._instance_decoder is not None:169 self._instance_decoder.reset_pooling_layer()170 if self._next_instance_decoder is not None:171 self._next_instance_decoder.reset_pooling_layer()172 173 def set_pool_size(self, pool_size):174 """Sets the pooling size of the ASPP pooling layers.175 176 Args:177 pool_size: A tuple specifying the pooling size of the ASPP pooling layers.178 """179 self._semantic_decoder.set_pool_size(pool_size)180 if self._instance_decoder is not None:181 self._instance_decoder.set_pool_size(pool_size)182 if self._next_instance_decoder is not None:183 self._next_instance_decoder.set_pool_size(pool_size)184 185 def get_pool_size(self):186 return self._semantic_decoder.get_pool_size()187 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 }211 items.update(instance_items)212 if self._next_instance_decoder is not None:213 next_instance_items = {214 common.CKPT_NEXT_INSTANCE_DECODER:215 self._next_instance_decoder,216 common.CKPT_NEXT_INSTANCE_REGRESSION_HEAD_WITHOUT_LAST_LAYER:217 self._next_instance_regression_head.conv_block,218 common.CKPT_NEXT_INSTANCE_REGRESSION_HEAD_LAST_LAYER:219 self._next_instance_regression_head.final_conv,220 }221 items.update(next_instance_items)222 return items223 224 def call(self, features, next_features, training=False):225 """Performs a forward pass.226 227 Args:228 features: An input dict of tf.Tensor with shape [batch, height, width,229 channels]. Different keys should point to different features extracted230 by the encoder, e.g. low-level or high-level features.231 next_features: An input dict of tf.Tensor similar to features. The232 features are computed with the next frame as input.233 training: A boolean flag indicating whether training behavior should be234 used (default: False).235 236 Returns:237 A dictionary containing the results of the semantic segmentation head and238 depending on the configuration also of the instance segmentation head.239 """240 241 semantic_features = self._semantic_decoder(features, training=training)242 results = self._semantic_head(semantic_features, training=training)243 244 if self._instance_decoder is not None:245 instance_features = self._instance_decoder(features, training=training)246 instance_center_predictions = self._instance_center_head(247 instance_features, training=training)248 instance_regression_predictions = self._instance_regression_head(249 instance_features, training=training)250 251 if results.keys() & instance_center_predictions.keys():252 raise ValueError('The keys of the semantic branch and the instance '253 'center branch overlap. Please use unique keys.')254 results.update(instance_center_predictions)255 256 if results.keys() & instance_regression_predictions.keys():257 raise ValueError('The keys of the semantic branch and the instance '258 'regression branch overlap. Please use unique keys.')259 results.update(instance_regression_predictions)260 261 if self._next_instance_decoder is not None:262 # We update the high level features in next_features with the concated263 # features of the high level features in both features and next_features.264 high_level_feature_name = self._next_high_level_feature_name265 high_level_features = features[high_level_feature_name]266 next_high_level_features = next_features[high_level_feature_name]267 next_high_level_features = tf.concat(268 [high_level_features, next_high_level_features], axis=3)269 next_features[high_level_feature_name] = next_high_level_features270 next_regression_features = self._next_instance_decoder(271 next_features, training=training)272 next_regression_predictions = self._next_instance_regression_head(273 next_regression_features, training=training)274 if results.keys() & next_regression_predictions.keys():275 raise ValueError('The keys of the next regresion branch overlap.'276 'Please use unique keys.')277 results.update(next_regression_predictions)278 279 return results280