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"""Tests for panoptic_deeplab."""17 18import tensorflow as tf19 20from deeplab2 import common21from deeplab2 import config_pb222from deeplab2.model.decoder import panoptic_deeplab23from deeplab2.utils import test_utils24 25 26def _create_panoptic_deeplab_example_proto(num_classes=19):27 semantic_decoder = config_pb2.DecoderOptions(28 feature_key='res5', atrous_rates=[6, 12, 18])29 semantic_head = config_pb2.HeadOptions(30 output_channels=num_classes, head_channels=256)31 32 instance_decoder = config_pb2.DecoderOptions(33 feature_key='res5', decoder_channels=128, atrous_rates=[6, 12, 18])34 center_head = config_pb2.HeadOptions(35 output_channels=1, head_channels=32)36 regression_head = config_pb2.HeadOptions(37 output_channels=2, head_channels=32)38 39 instance_branch = config_pb2.InstanceOptions(40 instance_decoder_override=instance_decoder,41 center_head=center_head,42 regression_head=regression_head)43 44 panoptic_deeplab_options = config_pb2.ModelOptions.PanopticDeeplabOptions(45 semantic_head=semantic_head, instance=instance_branch)46 # Add features from lowest to highest.47 panoptic_deeplab_options.low_level.add(48 feature_key='res3', channels_project=64)49 panoptic_deeplab_options.low_level.add(50 feature_key='res2', channels_project=32)51 52 return config_pb2.ModelOptions(53 decoder=semantic_decoder, panoptic_deeplab=panoptic_deeplab_options)54 55 56def _create_expected_shape(input_shape, output_channels):57 output_shape = input_shape.copy()58 output_shape[3] = output_channels59 return output_shape60 61 62class PanopticDeeplabTest(tf.test.TestCase):63 64 def test_panoptic_deeplab_single_decoder_init_errors(self):65 with self.assertRaises(ValueError):66 _ = panoptic_deeplab.PanopticDeepLabSingleDecoder(67 high_level_feature_name='test',68 low_level_feature_names=['only_one_name'], # Error: Only one name.69 low_level_channels_project=[64, 32],70 aspp_output_channels=256,71 decoder_output_channels=256,72 atrous_rates=[6, 12, 18],73 name='test_decoder')74 75 with self.assertRaises(ValueError):76 _ = panoptic_deeplab.PanopticDeepLabSingleDecoder(77 high_level_feature_name='test',78 low_level_feature_names=['one', 'two'],79 low_level_channels_project=[64], # Error: Only one projection size.80 aspp_output_channels=256,81 decoder_output_channels=256,82 atrous_rates=[6, 12, 18],83 name='test_decoder')84 85 def test_panoptic_deeplab_single_decoder_call_errors(self):86 decoder = panoptic_deeplab.PanopticDeepLabSingleDecoder(87 high_level_feature_name='high',88 low_level_feature_names=['low_one', 'low_two'],89 low_level_channels_project=[64, 32],90 aspp_output_channels=256,91 decoder_output_channels=256,92 atrous_rates=[6, 12, 18],93 name='test_decoder')94 95 with self.assertRaises(KeyError):96 input_dict = {'not_high': tf.random.uniform(shape=(2, 32, 32, 512)),97 'low_one': tf.random.uniform(shape=(2, 128, 128, 128)),98 'low_two': tf.random.uniform(shape=(2, 256, 256, 64))}99 _ = decoder(input_dict)100 with self.assertRaises(KeyError):101 input_dict = {'high': tf.random.uniform(shape=(2, 32, 32, 512)),102 'not_low_one': tf.random.uniform(shape=(2, 128, 128, 128)),103 'low_two': tf.random.uniform(shape=(2, 256, 256, 64))}104 _ = decoder(input_dict)105 with self.assertRaises(KeyError):106 input_dict = {'high': tf.random.uniform(shape=(2, 32, 32, 512)),107 'low_one': tf.random.uniform(shape=(2, 128, 128, 128)),108 'not_low_two': tf.random.uniform(shape=(2, 256, 256, 64))}109 _ = decoder(input_dict)110 111 def test_panoptic_deeplab_single_decoder_reset_pooling(self):112 decoder = panoptic_deeplab.PanopticDeepLabSingleDecoder(113 high_level_feature_name='high',114 low_level_feature_names=['low_one', 'low_two'],115 low_level_channels_project=[64, 32],116 aspp_output_channels=256,117 decoder_output_channels=256,118 atrous_rates=[6, 12, 18],119 name='test_decoder')120 pool_size = (None, None)121 decoder.reset_pooling_layer()122 123 self.assertTupleEqual(decoder._aspp._aspp_pool._pool_size,124 pool_size)125 126 def test_panoptic_deeplab_single_decoder_set_pooling(self):127 decoder = panoptic_deeplab.PanopticDeepLabSingleDecoder(128 high_level_feature_name='high',129 low_level_feature_names=['low_one', 'low_two'],130 low_level_channels_project=[64, 32],131 aspp_output_channels=256,132 decoder_output_channels=256,133 atrous_rates=[6, 12, 18],134 name='test_decoder')135 136 pool_size = (10, 10)137 decoder.set_pool_size(pool_size)138 139 self.assertTupleEqual(decoder._aspp._aspp_pool._pool_size,140 pool_size)141 142 def test_panoptic_deeplab_single_decoder_output_shape(self):143 decoder_channels = 256144 decoder = panoptic_deeplab.PanopticDeepLabSingleDecoder(145 high_level_feature_name='high',146 low_level_feature_names=['low_one', 'low_two'],147 low_level_channels_project=[64, 32],148 aspp_output_channels=256,149 decoder_output_channels=decoder_channels,150 atrous_rates=[6, 12, 18],151 name='test_decoder')152 153 input_shapes_list = [[[2, 128, 128, 128], [2, 256, 256, 64],154 [2, 32, 32, 512]],155 [[2, 129, 129, 128], [2, 257, 257, 64],156 [2, 33, 33, 512]]]157 158 for shapes in input_shapes_list:159 input_dict = {'low_one': tf.random.uniform(shape=shapes[0]),160 'low_two': tf.random.uniform(shape=shapes[1]),161 'high': tf.random.uniform(shape=shapes[2])}162 163 expected_shape = _create_expected_shape(shapes[1], decoder_channels)164 165 resulting_tensor = decoder(input_dict)166 self.assertListEqual(resulting_tensor.shape.as_list(), expected_shape)167 168 def test_panoptic_deeplab_single_head_output_shape(self):169 output_channels = 19170 head = panoptic_deeplab.PanopticDeepLabSingleHead(171 intermediate_channels=256,172 output_channels=output_channels,173 pred_key='pred',174 name='test_head')175 176 input_shapes_list = [[2, 256, 256, 48], [2, 257, 257, 48]]177 for shape in input_shapes_list:178 input_tensor = tf.random.uniform(shape=shape)179 expected_shape = _create_expected_shape(shape, output_channels)180 181 resulting_tensor = head(input_tensor)182 self.assertListEqual(resulting_tensor['pred'].shape.as_list(),183 expected_shape)184 185 def test_panoptic_deeplab_decoder_output_shape(self):186 num_classes = 31187 model_options = _create_panoptic_deeplab_example_proto(188 num_classes=num_classes)189 decoder = panoptic_deeplab.PanopticDeepLab(190 panoptic_deeplab_options=model_options.panoptic_deeplab,191 decoder_options=model_options.decoder)192 193 input_shapes_list = [[[2, 256, 256, 64], [2, 128, 128, 128],194 [2, 32, 32, 512]],195 [[2, 257, 257, 64], [2, 129, 129, 128],196 [2, 33, 33, 512]]]197 198 for shapes in input_shapes_list:199 input_dict = {'res2': tf.random.uniform(shape=shapes[0]),200 'res3': tf.random.uniform(shape=shapes[1]),201 'res5': tf.random.uniform(shape=shapes[2])}202 203 expected_semantic_shape = _create_expected_shape(shapes[0], num_classes)204 expected_instance_center_shape = _create_expected_shape(shapes[0], 1)205 expected_instance_regression_shape = _create_expected_shape(shapes[0], 2)206 207 resulting_dict = decoder(input_dict)208 self.assertListEqual(209 resulting_dict[common.PRED_SEMANTIC_LOGITS_KEY].shape.as_list(),210 expected_semantic_shape)211 self.assertListEqual(212 resulting_dict[common.PRED_CENTER_HEATMAP_KEY].shape.as_list(),213 expected_instance_center_shape)214 self.assertListEqual(215 resulting_dict[common.PRED_OFFSET_MAP_KEY].shape.as_list(),216 expected_instance_regression_shape)217 218 @test_utils.test_all_strategies219 def test_panoptic_deeplab_sync_bn(self, strategy):220 num_classes = 31221 model_options = _create_panoptic_deeplab_example_proto(222 num_classes=num_classes)223 input_dict = {'res2': tf.random.uniform(shape=[2, 257, 257, 64]),224 'res3': tf.random.uniform(shape=[2, 129, 129, 128]),225 'res5': tf.random.uniform(shape=[2, 33, 33, 512])}226 227 with strategy.scope():228 for bn_layer in test_utils.NORMALIZATION_LAYERS:229 decoder = panoptic_deeplab.PanopticDeepLab(230 panoptic_deeplab_options=model_options.panoptic_deeplab,231 decoder_options=model_options.decoder,232 bn_layer=bn_layer)233 _ = decoder(input_dict)234 235 def test_panoptic_deeplab_single_decoder_logging_feature_order(self):236 with self.assertLogs(level='WARN'):237 _ = panoptic_deeplab.PanopticDeepLabSingleDecoder(238 high_level_feature_name='high',239 low_level_feature_names=['low_two', 'low_one'],240 low_level_channels_project=[32, 64], # Potentially wrong order.241 aspp_output_channels=256,242 decoder_output_channels=256,243 atrous_rates=[6, 12, 18],244 name='test_decoder')245 246 def test_panoptic_deeplab_decoder_ckpt_tems(self):247 num_classes = 31248 model_options = _create_panoptic_deeplab_example_proto(249 num_classes=num_classes)250 decoder = panoptic_deeplab.PanopticDeepLab(251 panoptic_deeplab_options=model_options.panoptic_deeplab,252 decoder_options=model_options.decoder)253 ckpt_dict = decoder.checkpoint_items254 self.assertIn(common.CKPT_SEMANTIC_DECODER, ckpt_dict)255 self.assertIn(common.CKPT_SEMANTIC_HEAD_WITHOUT_LAST_LAYER, ckpt_dict)256 self.assertIn(common.CKPT_SEMANTIC_LAST_LAYER, ckpt_dict)257 self.assertIn(common.CKPT_INSTANCE_DECODER, ckpt_dict)258 self.assertIn(common.CKPT_INSTANCE_REGRESSION_HEAD_WITHOUT_LAST_LAYER,259 ckpt_dict)260 self.assertIn(common.CKPT_INSTANCE_REGRESSION_HEAD_LAST_LAYER, ckpt_dict)261 self.assertIn(common.CKPT_INSTANCE_CENTER_HEAD_WITHOUT_LAST_LAYER,262 ckpt_dict)263 self.assertIn(common.CKPT_INSTANCE_CENTER_HEAD_LAST_LAYER, ckpt_dict)264 265 266if __name__ == '__main__':267 tf.test.main()268