# Copyright 2018 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 mobilenet_v1.py.

This test mainly focuses on comparing slim MobilenetV1 and Keras MobilenetV1 for
object detection. To verify the consistency of the two models, we compare:
  1. Output shape of each layer given different inputs
  2. Number of global variables

We also visualize the model structure via Tensorboard, and compare the model
layout and the parameters of each Op to make sure the two implementations are
consistent.
"""

import itertools
import numpy as np
import tensorflow as tf

from google.protobuf import text_format

from object_detection.builders import hyperparams_builder
from object_detection.models.keras_models import mobilenet_v1
from object_detection.models.keras_models import test_utils
from object_detection.protos import hyperparams_pb2
from object_detection.utils import test_case

_KERAS_LAYERS_TO_CHECK = [
    'conv1_relu',
    'conv_dw_1_relu', 'conv_pw_1_relu',
    'conv_dw_2_relu', 'conv_pw_2_relu',
    'conv_dw_3_relu', 'conv_pw_3_relu',
    'conv_dw_4_relu', 'conv_pw_4_relu',
    'conv_dw_5_relu', 'conv_pw_5_relu',
    'conv_dw_6_relu', 'conv_pw_6_relu',
    'conv_dw_7_relu', 'conv_pw_7_relu',
    'conv_dw_8_relu', 'conv_pw_8_relu',
    'conv_dw_9_relu', 'conv_pw_9_relu',
    'conv_dw_10_relu', 'conv_pw_10_relu',
    'conv_dw_11_relu', 'conv_pw_11_relu',
    'conv_dw_12_relu', 'conv_pw_12_relu',
    'conv_dw_13_relu', 'conv_pw_13_relu',
]

_NUM_CHANNELS = 3
_BATCH_SIZE = 2


class MobilenetV1Test(test_case.TestCase):

  def _build_conv_hyperparams(self):
    conv_hyperparams = hyperparams_pb2.Hyperparams()
    conv_hyperparams_text_proto = """
      activation: RELU_6
      regularizer {
        l2_regularizer {
        }
      }
      initializer {
        truncated_normal_initializer {
        }
      }
      batch_norm {
        train: true,
        scale: false,
        center: true,
        decay: 0.2,
        epsilon: 0.1,
      }
    """
    text_format.Merge(conv_hyperparams_text_proto, conv_hyperparams)
    return hyperparams_builder.KerasLayerHyperparams(conv_hyperparams)

  def _create_application_with_layer_outputs(
      self, layer_names, batchnorm_training,
      conv_hyperparams=None,
      use_explicit_padding=False,
      alpha=1.0,
      min_depth=None):