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"""Tests for max_deeplab."""17 18import tensorflow as tf19 20from deeplab2 import common21from deeplab2 import config_pb222from deeplab2.model.decoder import max_deeplab23 24 25def _create_max_deeplab_example_proto(num_non_void_classes=19):26  semantic_decoder = config_pb2.DecoderOptions(27      feature_key='feature_semantic', atrous_rates=[6, 12, 18])28  auxiliary_semantic_head = config_pb2.HeadOptions(29      output_channels=num_non_void_classes, head_channels=256)30  pixel_space_head = config_pb2.HeadOptions(31      output_channels=128, head_channels=256)32  max_deeplab_options = config_pb2.ModelOptions.MaXDeepLabOptions(33      pixel_space_head=pixel_space_head,34      auxiliary_semantic_head=auxiliary_semantic_head)35  # Add features from lowest to highest.36  max_deeplab_options.auxiliary_low_level.add(37      feature_key='res3', channels_project=64)38  max_deeplab_options.auxiliary_low_level.add(39      feature_key='res2', channels_project=32)40  return config_pb2.ModelOptions(41      decoder=semantic_decoder, max_deeplab=max_deeplab_options)42 43 44class MaXDeeplabTest(tf.test.TestCase):45 46  def test_max_deeplab_decoder_output_shape(self):47    num_non_void_classes = 1948    num_mask_slots = 12749    model_options = _create_max_deeplab_example_proto(50        num_non_void_classes=num_non_void_classes)51    decoder = max_deeplab.MaXDeepLab(52        max_deeplab_options=model_options.max_deeplab,53        ignore_label=255,54        decoder_options=model_options.decoder)55 56    input_dict = {57        'res2':58            tf.random.uniform([2, 17, 17, 256]),59        'res3':60            tf.random.uniform([2, 9, 9, 512]),61        'transformer_class_feature':62            tf.random.uniform([2, num_mask_slots, 256]),63        'transformer_mask_feature':64            tf.random.uniform([2, num_mask_slots, 256]),65        'feature_panoptic':66            tf.random.uniform([2, 17, 17, 256]),67        'feature_semantic':68            tf.random.uniform([2, 5, 5, 2048])69    }70    resulting_dict = decoder(input_dict)71    self.assertListEqual(72        resulting_dict[common.PRED_SEMANTIC_LOGITS_KEY].shape.as_list(),73        [2, 17, 17, 19])  # Stride 474    self.assertListEqual(75        resulting_dict[76            common.PRED_PIXEL_SPACE_NORMALIZED_FEATURE_KEY].shape.as_list(),77        [2, 17, 17, 128])  # Stride 478    self.assertListEqual(79        resulting_dict[80            common.PRED_TRANSFORMER_CLASS_LOGITS_KEY].shape.as_list(),81        # Non-void classes and a void class.82        [2, num_mask_slots, num_non_void_classes + 1])83    self.assertListEqual(84        resulting_dict[common.PRED_PIXEL_SPACE_MASK_LOGITS_KEY].shape.as_list(),85        [2, 17, 17, num_mask_slots])  # Stride 4.86 87 88if __name__ == '__main__':89  tf.test.main()90