| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| """This file contains the code for the Motion-DeepLab decoder.""" |
|
|
| import tensorflow as tf |
|
|
| from deeplab2 import common |
| from deeplab2 import config_pb2 |
| from deeplab2.model.decoder import panoptic_deeplab |
|
|
|
|
| class MotionDeepLabDecoder(tf.keras.layers.Layer): |
| """A Motion-DeepLab decoder layer. |
| |
| This layer takes low- and high-level features as input and uses a dual-ASPP |
| and dual-decoder structure to aggregate features for semantic and instance |
| segmentation. On top of the decoders, four heads are used to predict semantic |
| segmentation, instance center probabilities, instance center regression, and |
| previous frame offset regression per pixel. |
| """ |
|
|
| def __init__( |
| self, |
| decoder_options: config_pb2.DecoderOptions, |
| motion_deeplab_options: config_pb2.ModelOptions.MotionDeepLabOptions, |
| bn_layer=tf.keras.layers.BatchNormalization): |
| """Initializes a Motion-DeepLab decoder. |
| |
| Args: |
| decoder_options: Decoder options as defined in config_pb2.DecoderOptions. |
| motion_deeplab_options: Model options as defined in |
| config_pb2.ModelOptions.MotionDeeplabOptions. |
| bn_layer: An optional tf.keras.layers.Layer that computes the |
| normalization (default: tf.keras.layers.BatchNormalization). |
| """ |
| super(MotionDeepLabDecoder, self).__init__(name='MotionDeepLabDecoder') |
|
|
| low_level_feature_keys = [ |
| item.feature_key for item in motion_deeplab_options.low_level |
| ] |
| low_level_channels_project = [ |
| item.channels_project for item in motion_deeplab_options.low_level |
| ] |
|
|
| self._semantic_decoder = panoptic_deeplab.PanopticDeepLabSingleDecoder( |
| decoder_options.feature_key, |
| low_level_feature_keys, |
| low_level_channels_project, |
| decoder_options.aspp_channels, |
| decoder_options.decoder_channels, |
| decoder_options.atrous_rates, |
| name='semantic_decoder', |
| bn_layer=bn_layer) |
| self._semantic_head = panoptic_deeplab.PanopticDeepLabSingleHead( |
| motion_deeplab_options.semantic_head.head_channels, |
| motion_deeplab_options.semantic_head.output_channels, |
| common.PRED_SEMANTIC_LOGITS_KEY, |
| name='semantic_head', |
| bn_layer=bn_layer) |
|
|
| self._instance_decoder = None |
| self._instance_center_head = None |
| self._instance_regression_head = None |
| self._motion_regression_head = None |
|
|
| if motion_deeplab_options.instance.low_level_override: |
| low_level_options = motion_deeplab_options.instance.low_level_override |
| else: |
| low_level_options = motion_deeplab_options.low_level |
|
|
| |
| |
| if motion_deeplab_options.instance.HasField('instance_decoder_override'): |
| decoder_options = (motion_deeplab_options.instance |
| .instance_decoder_override) |
|
|
| low_level_feature_keys = [item.feature_key for item in low_level_options] |
| low_level_channels_project = [ |
| item.channels_project for item in low_level_options |
| ] |
|
|
| self._instance_decoder = panoptic_deeplab.PanopticDeepLabSingleDecoder( |
| decoder_options.feature_key, |
| low_level_feature_keys, |
| low_level_channels_project, |
| decoder_options.aspp_channels, |
| decoder_options.decoder_channels, |
| decoder_options.atrous_rates, |
| name='instance_decoder', |
| bn_layer=bn_layer) |
| self._instance_center_head = panoptic_deeplab.PanopticDeepLabSingleHead( |
| motion_deeplab_options.instance.center_head.head_channels, |
| motion_deeplab_options.instance.center_head.output_channels, |
| common.PRED_CENTER_HEATMAP_KEY, |
| name='instance_center_head', |
| bn_layer=bn_layer) |
| self._instance_regression_head = panoptic_deeplab.PanopticDeepLabSingleHead( |
| motion_deeplab_options.instance.regression_head.head_channels, |
| motion_deeplab_options.instance.regression_head.output_channels, |
| common.PRED_OFFSET_MAP_KEY, |
| name='instance_regression_head', |
| bn_layer=bn_layer) |
|
|
| |
| |
| self._motion_regression_head = panoptic_deeplab.PanopticDeepLabSingleHead( |
| motion_deeplab_options.motion_head.head_channels, |
| motion_deeplab_options.motion_head.output_channels, |
| common.PRED_FRAME_OFFSET_MAP_KEY, |
| name='motion_regression_head', |
| bn_layer=bn_layer) |
|
|
| def reset_pooling_layer(self): |
| """Resets the ASPP pooling layers to global average pooling.""" |
| self._semantic_decoder.reset_pooling_layer() |
| if self._instance_decoder is not None: |
| self._instance_decoder.reset_pooling_layer() |
|
|
| def set_pool_size(self, pool_size): |
| """Sets the pooling size of the ASPP pooling layers. |
| |
| Args: |
| pool_size: A tuple specifying the pooling size of the ASPP pooling layers. |
| """ |
| self._semantic_decoder.set_pool_size(pool_size) |
| if self._instance_decoder is not None: |
| self._instance_decoder.set_pool_size(pool_size) |
|
|
| def get_pool_size(self): |
| return self._semantic_decoder.get_pool_size() |
|
|
| def call(self, features, training=False): |
| """Performs a forward pass. |
| |
| Args: |
| features: An input dict of tf.Tensor with shape [batch, height, width, |
| channels]. Different keys should point to different features extracted |
| by the encoder, e.g. low-level or high-level features. |
| training: A boolean flag indicating whether training behavior should be |
| used (default: False). |
| |
| Returns: |
| A dictionary containing the results of the semantic segmentation head and |
| depending on the configuration also of the instance segmentation head. |
| """ |
|
|
| semantic_features = self._semantic_decoder(features, training=training) |
| results = self._semantic_head(semantic_features, training=training) |
|
|
| if self._instance_decoder is not None: |
| instance_features = self._instance_decoder(features, training=training) |
| instance_center_predictions = self._instance_center_head( |
| instance_features, training=training) |
| instance_regression_predictions = self._instance_regression_head( |
| instance_features, training=training) |
| motion_regression_predictions = self._motion_regression_head( |
| instance_features, training=training) |
| if results.keys() & motion_regression_predictions.keys(): |
| raise ValueError('The keys of the semantic branch and the instance ' |
| 'motion branch overlap. Please use unique keys.') |
| results.update(motion_regression_predictions) |
|
|
| if results.keys() & instance_center_predictions.keys(): |
| raise ValueError('The keys of the semantic branch and the instance ' |
| 'center branch overlap. Please use unique keys.') |
| results.update(instance_center_predictions) |
|
|
| if results.keys() & instance_regression_predictions.keys(): |
| raise ValueError('The keys of the semantic branch and the instance ' |
| 'regression branch overlap. Please use unique keys.') |
| results.update(instance_regression_predictions) |
|
|
| return results |
|
|
| @property |
| def checkpoint_items(self): |
| items = { |
| common.CKPT_SEMANTIC_DECODER: |
| self._semantic_decoder, |
| common.CKPT_SEMANTIC_HEAD_WITHOUT_LAST_LAYER: |
| self._semantic_head.conv_block, |
| common.CKPT_SEMANTIC_LAST_LAYER: |
| self._semantic_head.final_conv |
| } |
| if self._instance_decoder is not None: |
| instance_items = { |
| common.CKPT_INSTANCE_DECODER: |
| self._instance_decoder, |
| common.CKPT_INSTANCE_CENTER_HEAD_WITHOUT_LAST_LAYER: |
| self._instance_center_head.conv_block, |
| common.CKPT_INSTANCE_CENTER_HEAD_LAST_LAYER: |
| self._instance_center_head.final_conv, |
| common.CKPT_INSTANCE_REGRESSION_HEAD_WITHOUT_LAST_LAYER: |
| self._instance_regression_head.conv_block, |
| common.CKPT_INSTANCE_REGRESSION_HEAD_LAST_LAYER: |
| self._instance_regression_head.final_conv, |
| common.CKPT_MOTION_REGRESSION_HEAD_WITHOUT_LAST_LAYER: |
| self._motion_regression_head.conv_block, |
| common.CKPT_MOTION_REGRESSION_HEAD_LAST_LAYER: |
| self._motion_regression_head.final_conv, |
| } |
| items.update(instance_items) |
| return items |
|
|
