# Lint as: python3
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for factory functions."""
# Import libraries
from absl.testing import parameterized
import tensorflow as tf

from tensorflow.python.distribute import combinations
from official.vision.beta.configs import backbones as backbones_cfg
from official.vision.beta.configs import backbones_3d as backbones_3d_cfg
from official.vision.beta.configs import common as common_cfg
from official.vision.beta.configs import retinanet as retinanet_cfg
from official.vision.beta.modeling import backbones
from official.vision.beta.modeling.backbones import factory


class FactoryTest(tf.test.TestCase, parameterized.TestCase):

  @combinations.generate(
      combinations.combine(model_id=[18, 34, 50, 101, 152],))
  def test_resnet_creation(self, model_id):
    """Test creation of ResNet models."""

    network = backbones.ResNet(
        model_id=model_id, norm_momentum=0.99, norm_epsilon=1e-5)

    backbone_config = backbones_cfg.Backbone(
        type='resnet',
        resnet=backbones_cfg.ResNet(model_id=model_id))
    norm_activation_config = common_cfg.NormActivation(
        norm_momentum=0.99, norm_epsilon=1e-5)
    model_config = retinanet_cfg.RetinaNet(
        backbone=backbone_config, norm_activation=norm_activation_config)

    factory_network = factory.build_backbone(
        input_specs=tf.keras.layers.InputSpec(shape=[None, None, None, 3]),
        model_config=model_config)

    network_config = network.get_config()
    factory_network_config = factory_network.get_config()

    self.assertEqual(network_config, factory_network_config)

  @combinations.generate(
      combinations.combine(
          model_id=['b0', 'b1', 'b2', 'b3', 'b4', 'b5', 'b6', 'b7'],
          se_ratio=[0.0, 0.25],
      ))
  def test_efficientnet_creation(self, model_id, se_ratio):
    """Test creation of EfficientNet models."""

    network = backbones.EfficientNet(
        model_id=model_id,
        se_ratio=se_ratio,
        norm_momentum=0.99,
        norm_epsilon=1e-5)

    backbone_config = backbones_cfg.Backbone(
        type='efficientnet',
        efficientnet=backbones_cfg.EfficientNet(
            model_id=model_id, se_ratio=se_ratio))
    norm_activation_config = common_cfg.NormActivation(
        norm_momentum=0.99, norm_epsilon=1e-5)
    model_config = retinanet_cfg.RetinaNet(
        backbone=backbone_config, norm_activation=norm_activation_config)

    factory_network = factory.build_backbone(
        input_specs=tf.keras.layers.InputSpec(shape=[None, None, None, 3]),
        model_config=model_config)

    network_config = network.get_config()
    factory_network_config = factory_network.get_config()

    self.assertEqual(network_config, factory_network_config)

  @combinations.generate(combinations.combine(model_id=['49'],))
  def test_spinenet_creation(self, model_id):
    """Test creation of SpineNet models."""
    input_size = 128
    min_level = 3
    max_level = 7

    input_specs = tf.keras.layers.InputSpec(
        shape=[None, input_size, input_size, 3])
    network = backbones.SpineNet(
        input_specs=input_specs,
        min_level=min_level,
        max_level=max_level,
        norm_momentum=0.99,
        norm_epsilon=1e-5)

    backbone_config = backbones_cfg.Backbone(
        type='spinenet',
        spinenet=backbones_cfg.SpineNet(model_id=model_id))
    norm_activation_config = common_cfg.NormActivation(
        norm_momentum=0.99, norm_epsilon=1e-5)
    model_config = retinanet_cfg.RetinaNet(
        backbone=backbone_config, norm_activation=norm_activation_config)

    factory_network = factory.build_backbone(
        input_specs=tf.keras.layers.InputSpec(
            shape=[None, input_size, input_size, 3]),
        model_config=model_config)

    network_config = network.get_config()
    factory_network_config = factory_network.get_config()

    self.assertEqual(network_config, factory_network_config)

  @combinations.generate(
      combinations.combine(model_id=[38, 56, 104],))
  def test_revnet_creation(self, model_id):
    """Test creation of RevNet models."""
    network = backbones.RevNet(
        model_id=model_id, norm_momentum=0.99, norm_epsilon=1e-5)

    backbone_config = backbones_cfg.Backbone(
        type='revnet',
        revnet=backbones_cfg.RevNet(model_id=model_id))
    norm_activation_config = common_cfg.NormActivation(
        norm_momentum=0.99, norm_epsilon=1e-5)
    model_config = retinanet_cfg.RetinaNet(
        backbone=backbone_config, norm_activation=norm_activation_config)

    factory_network = factory.build_backbone(
        input_specs=tf.keras.layers.InputSpec(shape=[None, None, None, 3]),
        model_config=model_config)

    network_config = network.get_config()
    factory_network_config = factory_network.get_config()

    self.assertEqual(network_config, factory_network_config)

  @combinations.generate(combinations.combine(model_type=['resnet_3d'],))
  def test_resnet_3d_creation(self, model_type):
    """Test creation of ResNet 3D models."""
    backbone_cfg = backbones_3d_cfg.Backbone3D(type=model_type).get()
    temporal_strides = []
    temporal_kernel_sizes = []
    for block_spec in backbone_cfg.block_specs:
      temporal_strides.append(block_spec.temporal_strides)
      temporal_kernel_sizes.append(block_spec.temporal_kernel_sizes)

    _ = backbones.ResNet3D(
        model_id=backbone_cfg.model_id,
        temporal_strides=temporal_strides,
        temporal_kernel_sizes=temporal_kernel_sizes,
        norm_momentum=0.99,
        norm_epsilon=1e-5)


if __name__ == '__main__':
  tf.test.main()
