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 model.builder."""17 18import os19from absl.testing import parameterized20 21import tensorflow as tf22 23from google.protobuf import text_format24from deeplab2 import config_pb225from deeplab2.model import builder26from deeplab2.model.decoder import motion_deeplab_decoder27from deeplab2.model.encoder import axial_resnet_instances28from deeplab2.model.encoder import mobilenet29# resources dependency30 31 32_CONFIG_PATH = 'deeplab2/configs/example'33 34 35def _read_proto_file(filename, proto):36  filename = filename  # OSS: removed internal filename loading.37  with tf.io.gfile.GFile(filename, 'r') as proto_file:38    return text_format.ParseLines(proto_file, proto)39 40 41class BuilderTest(tf.test.TestCase, parameterized.TestCase):42 43  def test_resnet50_encoder_creation(self):44    backbone_options = config_pb2.ModelOptions.BackboneOptions(45        name='resnet50', output_stride=32)46    encoder = builder.create_encoder(47        backbone_options,48        tf.keras.layers.experimental.SyncBatchNormalization)49    self.assertIsInstance(encoder, axial_resnet_instances.ResNet50)50 51  @parameterized.parameters('mobilenet_v3_large', 'mobilenet_v3_small')52  def test_mobilenet_encoder_creation(self, model_name):53    backbone_options = config_pb2.ModelOptions.BackboneOptions(54        name=model_name, use_squeeze_and_excite=True, output_stride=32)55    encoder = builder.create_encoder(56        backbone_options,57        tf.keras.layers.experimental.SyncBatchNormalization)58    self.assertIsInstance(encoder, mobilenet.MobileNet)59 60  def test_resnet_encoder_creation(self):61    backbone_options = config_pb2.ModelOptions.BackboneOptions(62        name='max_deeplab_s', output_stride=32)63    encoder = builder.create_resnet_encoder(64        backbone_options,65        bn_layer=tf.keras.layers.experimental.SyncBatchNormalization)66    self.assertIsInstance(encoder, axial_resnet_instances.MaXDeepLabS)67 68  def test_decoder_creation(self):69    proto_filename = os.path.join(70        _CONFIG_PATH, 'example_kitti-step_motion_deeplab.textproto')71    model_options = _read_proto_file(proto_filename, config_pb2.ModelOptions())72    motion_decoder = builder.create_decoder(73        model_options, tf.keras.layers.experimental.SyncBatchNormalization,74        ignore_label=255)75    self.assertIsInstance(motion_decoder,76                          motion_deeplab_decoder.MotionDeepLabDecoder)77 78 79if __name__ == '__main__':80  tf.test.main()81