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 Panoptic-DeepLab decoder.17 18Reference:19 - [Panoptic-DeepLab: A Simple, Strong, and Fast Baseline for Bottom-Up20 Panoptic Segmentation](https://arxiv.org/pdf/1911.10194)21"""22from absl import logging23 24import tensorflow as tf25 26from deeplab2 import common27from deeplab2.model import utils28from deeplab2.model.decoder import aspp29from deeplab2.model.layers import convolutions30 31 32layers = tf.keras.layers33 34 35class PanopticDeepLabSingleDecoder(layers.Layer):36 """A single Panoptic-DeepLab decoder layer.37 38 This layer takes low- and high-level features as input and uses an ASPP39 followed by a fusion block to decode features for a single task, e.g.,40 semantic segmentation or instance segmentation.41 """42 43 def __init__(self,44 high_level_feature_name,45 low_level_feature_names,46 low_level_channels_project,47 aspp_output_channels,48 decoder_output_channels,49 atrous_rates,50 name,51 aspp_use_only_1x1_proj_conv=False,52 decoder_conv_type='depthwise_separable_conv',53 bn_layer=tf.keras.layers.BatchNormalization):54 """Initializes a single Panoptic-DeepLab decoder of layers.Layer.55 56 Args:57 high_level_feature_name: A string specifying the name of the high-level58 feature coming from an encoder.59 low_level_feature_names: A list of strings specifying the name of the60 low-level features coming from an encoder. An order from highest to61 lower level is expected, e.g. ['res3', 'res2'].62 low_level_channels_project: A list of integer specifying the number of63 filters used for processing each low_level features.64 aspp_output_channels: An integer specifying the number of filters in the65 ASPP convolution layers.66 decoder_output_channels: An integer specifying the number of filters in67 the decoder convolution layers.68 atrous_rates: A list of three integers specifying the atrous rate for the69 ASPP layers.70 name: A string specifying the name of the layer.71 aspp_use_only_1x1_proj_conv: Boolean, specifying if the ASPP five branches72 are turned off or not. If True, the ASPP module is degenerated to one73 1x1 convolution, projecting the input channels to `output_channels`.74 decoder_conv_type: String, specifying decoder convolution type. Support75 'depthwise_separable_conv' and 'standard_conv'.76 bn_layer: An optional tf.keras.layers.Layer that computes the77 normalization (default: tf.keras.layers.BatchNormalization).78 79 Raises:80 ValueError: An error occurs when the length of low_level_feature_names81 differs from the length of low_level_channels_project.82 """83 super(PanopticDeepLabSingleDecoder, self).__init__(name=name)84 self._channel_axis = 385 86 self._aspp = aspp.ASPP(87 aspp_output_channels,88 atrous_rates,89 aspp_use_only_1x1_proj_conv=aspp_use_only_1x1_proj_conv,90 name='aspp',91 bn_layer=bn_layer)92 self._high_level_feature_name = high_level_feature_name93 94 if len(low_level_feature_names) != len(low_level_channels_project):95 raise ValueError('The Panoptic-DeepLab decoder requires the same number '96 'of low-level features as the number of low-level '97 'projection channels. But got %d and %d.'98 % (len(low_level_feature_names),99 len(low_level_channels_project)))100 101 self._low_level_feature_names = low_level_feature_names102 103 for i, channels_project in enumerate(low_level_channels_project):104 # Check if channel sizes increases and issue a warning.105 if i > 0 and low_level_channels_project[i - 1] < channels_project:106 logging.warning(107 'The low level projection channels usually do not '108 'increase for features with higher spatial resolution. '109 'Please make sure, this behavior is intended.')110 current_low_level_conv_name, current_fusion_conv_name = (111 utils.get_low_level_conv_fusion_conv_current_names(i))112 utils.safe_setattr(113 self, current_low_level_conv_name, convolutions.Conv2DSame(114 channels_project,115 kernel_size=1,116 name=utils.get_layer_name(current_low_level_conv_name),117 use_bias=False,118 use_bn=True,119 bn_layer=bn_layer,120 activation='relu'))121 122 utils.safe_setattr(123 self, current_fusion_conv_name, convolutions.StackedConv2DSame(124 conv_type=decoder_conv_type,125 num_layers=1,126 output_channels=decoder_output_channels,127 kernel_size=5,128 name=utils.get_layer_name(current_fusion_conv_name),129 use_bias=False,130 use_bn=True,131 bn_layer=bn_layer,132 activation='relu'))133 134 def call(self, features, training=False):135 """Performs a forward pass.136 137 Args:138 features: An input dict of tf.Tensor with shape [batch, height, width,139 channels]. Different keys should point to different features extracted140 by the encoder, e.g. low-level or high-level features.141 training: A boolean flag indicating whether training behavior should be142 used (default: False).143 144 Returns:145 Refined features as instance of tf.Tensor.146 """147 148 high_level_features = features[self._high_level_feature_name]149 combined_features = self._aspp(high_level_features, training=training)150 151 # Fuse low-level features with high-level features.152 for i in range(len(self._low_level_feature_names)):153 current_low_level_conv_name, current_fusion_conv_name = (154 utils.get_low_level_conv_fusion_conv_current_names(i))155 # Iterate from the highest level of the low level features to the lowest156 # level, i.e. take the features with the smallest spatial size first.157 low_level_features = features[self._low_level_feature_names[i]]158 low_level_features = getattr(self, current_low_level_conv_name)(159 low_level_features, training=training)160 161 target_h = tf.shape(low_level_features)[1]162 target_w = tf.shape(low_level_features)[2]163 source_h = tf.shape(combined_features)[1]164 source_w = tf.shape(combined_features)[2]165 166 tf.assert_less(167 source_h - 1,168 target_h,169 message='Features are down-sampled during decoder.')170 tf.assert_less(171 source_w - 1,172 target_w,173 message='Features are down-sampled during decoder.')174 175 combined_features = utils.resize_align_corners(combined_features,176 [target_h, target_w])177 178 combined_features = tf.concat([combined_features, low_level_features],179 self._channel_axis)180 combined_features = getattr(self, current_fusion_conv_name)(181 combined_features, training=training)182 183 return combined_features184 185 def reset_pooling_layer(self):186 """Resets the ASPP pooling layer to global average pooling."""187 self._aspp.reset_pooling_layer()188 189 def set_pool_size(self, pool_size):190 """Sets the pooling size of the ASPP pooling layer.191 192 Args:193 pool_size: A tuple specifying the pooling size of the ASPP pooling layer.194 """195 self._aspp.set_pool_size(pool_size)196 197 def get_pool_size(self):198 return self._aspp.get_pool_size()199 200 201class PanopticDeepLabSingleHead(layers.Layer):202 """A single PanopticDeepLab head layer.203 204 This layer takes in the enriched features from a decoder and adds two205 convolutions on top.206 """207 208 def __init__(self,209 intermediate_channels,210 output_channels,211 pred_key,212 name,213 conv_type='depthwise_separable_conv',214 bn_layer=tf.keras.layers.BatchNormalization):215 """Initializes a single PanopticDeepLab head.216 217 Args:218 intermediate_channels: An integer specifying the number of filters of the219 first 5x5 convolution.220 output_channels: An integer specifying the number of filters of the second221 1x1 convolution.222 pred_key: A string specifying the key of the output dictionary.223 name: A string specifying the name of this head.224 conv_type: String, specifying head convolution type. Support225 'depthwise_separable_conv' and 'standard_conv'.226 bn_layer: An optional tf.keras.layers.Layer that computes the227 normalization (default: tf.keras.layers.BatchNormalization).228 """229 super(PanopticDeepLabSingleHead, self).__init__(name=name)230 self._pred_key = pred_key231 232 self.conv_block = convolutions.StackedConv2DSame(233 conv_type=conv_type,234 num_layers=1,235 output_channels=intermediate_channels,236 kernel_size=5,237 name='conv_block',238 use_bias=False,239 use_bn=True,240 bn_layer=bn_layer,241 activation='relu')242 self.final_conv = layers.Conv2D(243 output_channels,244 kernel_size=1,245 name='final_conv',246 kernel_initializer=tf.keras.initializers.TruncatedNormal(stddev=0.01))247 248 def call(self, features, training=False):249 """Performs a forward pass.250 251 Args:252 features: A tf.Tensor with shape [batch, height, width, channels].253 training: A boolean flag indicating whether training behavior should be254 used (default: False).255 256 Returns:257 The dictionary containing the predictions under the specified key.258 """259 x = self.conv_block(features, training=training)260 return {self._pred_key: self.final_conv(x)}261 262 263class PanopticDeepLab(layers.Layer):264 """A Panoptic-DeepLab decoder layer.265 266 This layer takes low- and high-level features as input and uses a dual-ASPP267 and dual-decoder structure to aggregate features for semantic and instance268 segmentation. On top of the decoders, three heads are used to predict semantic269 segmentation, instance center probabilities, and instance center regression270 per pixel.271 """272 273 def __init__(self,274 decoder_options,275 panoptic_deeplab_options,276 bn_layer=tf.keras.layers.BatchNormalization):277 """Initializes a Panoptic-DeepLab decoder.278 279 Args:280 decoder_options: Decoder options as defined in config_pb2.DecoderOptions.281 panoptic_deeplab_options: Model options as defined in282 config_pb2.ModelOptions.PanopticDeeplabOptions.283 bn_layer: An optional tf.keras.layers.Layer that computes the284 normalization (default: tf.keras.layers.BatchNormalization).285 """286 super(PanopticDeepLab, self).__init__(name='PanopticDeepLab')287 288 low_level_feature_keys = [289 item.feature_key for item in panoptic_deeplab_options.low_level290 ]291 low_level_channels_project = [292 item.channels_project for item in panoptic_deeplab_options.low_level293 ]294 295 self._semantic_decoder = PanopticDeepLabSingleDecoder(296 high_level_feature_name=decoder_options.feature_key,297 low_level_feature_names=low_level_feature_keys,298 low_level_channels_project=low_level_channels_project,299 aspp_output_channels=decoder_options.aspp_channels,300 decoder_output_channels=decoder_options.decoder_channels,301 atrous_rates=decoder_options.atrous_rates,302 name='semantic_decoder',303 aspp_use_only_1x1_proj_conv=decoder_options.aspp_use_only_1x1_proj_conv,304 decoder_conv_type=decoder_options.decoder_conv_type,305 bn_layer=bn_layer)306 self._semantic_head = PanopticDeepLabSingleHead(307 panoptic_deeplab_options.semantic_head.head_channels,308 panoptic_deeplab_options.semantic_head.output_channels,309 common.PRED_SEMANTIC_LOGITS_KEY,310 name='semantic_head',311 conv_type=panoptic_deeplab_options.semantic_head.head_conv_type,312 bn_layer=bn_layer)313 314 self._instance_decoder = None315 self._instance_center_head = None316 self._instance_regression_head = None317 318 if panoptic_deeplab_options.instance.enable:319 if panoptic_deeplab_options.instance.low_level_override:320 low_level_options = panoptic_deeplab_options.instance.low_level_override321 else:322 low_level_options = panoptic_deeplab_options.low_level323 324 # If instance_decoder is set, use those options; otherwise reuse the325 # architecture as defined for the semantic decoder.326 if panoptic_deeplab_options.instance.HasField(327 'instance_decoder_override'):328 decoder_options = (panoptic_deeplab_options.instance329 .instance_decoder_override)330 331 low_level_feature_keys = [item.feature_key for item in low_level_options]332 low_level_channels_project = [333 item.channels_project for item in low_level_options334 ]335 336 self._instance_decoder = PanopticDeepLabSingleDecoder(337 high_level_feature_name=decoder_options.feature_key,338 low_level_feature_names=low_level_feature_keys,339 low_level_channels_project=low_level_channels_project,340 aspp_output_channels=decoder_options.aspp_channels,341 decoder_output_channels=decoder_options.decoder_channels,342 atrous_rates=decoder_options.atrous_rates,343 name='instance_decoder',344 aspp_use_only_1x1_proj_conv=(345 decoder_options.aspp_use_only_1x1_proj_conv),346 decoder_conv_type=decoder_options.decoder_conv_type,347 bn_layer=bn_layer)348 self._instance_center_head = PanopticDeepLabSingleHead(349 panoptic_deeplab_options.instance.center_head.head_channels,350 panoptic_deeplab_options.instance.center_head.output_channels,351 common.PRED_CENTER_HEATMAP_KEY,352 name='instance_center_head',353 conv_type=(354 panoptic_deeplab_options.instance.center_head.head_conv_type),355 bn_layer=bn_layer)356 self._instance_regression_head = PanopticDeepLabSingleHead(357 panoptic_deeplab_options.instance.regression_head.head_channels,358 panoptic_deeplab_options.instance.regression_head.output_channels,359 common.PRED_OFFSET_MAP_KEY,360 name='instance_regression_head',361 conv_type=(362 panoptic_deeplab_options.instance.regression_head.head_conv_type),363 bn_layer=bn_layer)364 365 def reset_pooling_layer(self):366 """Resets the ASPP pooling layers to global average pooling."""367 self._semantic_decoder.reset_pooling_layer()368 if self._instance_decoder is not None:369 self._instance_decoder.reset_pooling_layer()370 371 def set_pool_size(self, pool_size):372 """Sets the pooling size of the ASPP pooling layers.373 374 Args:375 pool_size: A tuple specifying the pooling size of the ASPP pooling layers.376 """377 self._semantic_decoder.set_pool_size(pool_size)378 if self._instance_decoder is not None:379 self._instance_decoder.set_pool_size(pool_size)380 381 def get_pool_size(self):382 return self._semantic_decoder.get_pool_size()383 384 @property385 def checkpoint_items(self):386 items = {387 common.CKPT_SEMANTIC_DECODER:388 self._semantic_decoder,389 common.CKPT_SEMANTIC_HEAD_WITHOUT_LAST_LAYER:390 self._semantic_head.conv_block,391 common.CKPT_SEMANTIC_LAST_LAYER:392 self._semantic_head.final_conv393 }394 if self._instance_decoder is not None:395 instance_items = {396 common.CKPT_INSTANCE_DECODER:397 self._instance_decoder,398 common.CKPT_INSTANCE_CENTER_HEAD_WITHOUT_LAST_LAYER:399 self._instance_center_head.conv_block,400 common.CKPT_INSTANCE_CENTER_HEAD_LAST_LAYER:401 self._instance_center_head.final_conv,402 common.CKPT_INSTANCE_REGRESSION_HEAD_WITHOUT_LAST_LAYER:403 self._instance_regression_head.conv_block,404 common.CKPT_INSTANCE_REGRESSION_HEAD_LAST_LAYER:405 self._instance_regression_head.final_conv,406 }407 items.update(instance_items)408 return items409 410 def call(self, features, training=False):411 """Performs a forward pass.412 413 Args:414 features: An input dict of tf.Tensor with shape [batch, height, width,415 channels]. Different keys should point to different features extracted416 by the encoder, e.g. low-level or high-level features.417 training: A boolean flag indicating whether training behavior should be418 used (default: False).419 420 Returns:421 A dictionary containing the results of the semantic segmentation head and422 depending on the configuration also of the instance segmentation head.423 """424 425 semantic_features = self._semantic_decoder(features, training=training)426 results = self._semantic_head(semantic_features, training=training)427 428 if self._instance_decoder is not None:429 instance_features = self._instance_decoder(features, training=training)430 instance_center_predictions = self._instance_center_head(431 instance_features, training=training)432 instance_regression_predictions = self._instance_regression_head(433 instance_features, training=training)434 435 if results.keys() & instance_center_predictions.keys():436 raise ValueError('The keys of the semantic branch and the instance '437 'center branch overlap. Please use unique keys.')438 results.update(instance_center_predictions)439 440 if results.keys() & instance_regression_predictions.keys():441 raise ValueError('The keys of the semantic branch and the instance '442 'regression branch overlap. Please use unique keys.')443 results.update(instance_regression_predictions)444 445 return results446