# 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_v2."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import copy
import tensorflow as tf
from nets.mobilenet import conv_blocks as ops
from nets.mobilenet import mobilenet
from nets.mobilenet import mobilenet_v2


slim = tf.contrib.slim


def find_ops(optype):
  """Find ops of a given type in graphdef or a graph.

  Args:
    optype: operation type (e.g. Conv2D)
  Returns:
     List of operations.
  """
  gd = tf.get_default_graph()
  return [var for var in gd.get_operations() if var.type == optype]


class MobilenetV2Test(tf.test.TestCase):

  def setUp(self):
    tf.reset_default_graph()

  def testCreation(self):
    spec = dict(mobilenet_v2.V2_DEF)
    _, ep = mobilenet.mobilenet(
        tf.placeholder(tf.float32, (10, 224, 224, 16)), conv_defs=spec)
    num_convs = len(find_ops('Conv2D'))

    # This is mostly a sanity test. No deep reason for these particular
    # constants.
    #
    # All but first 2 and last one have  two convolutions, and there is one
    # extra conv that is not in the spec. (logits)
    self.assertEqual(num_convs, len(spec['spec']) * 2 - 2)
    # Check that depthwise are exposed.
    for i in range(2, 17):
      self.assertIn('layer_%d/depthwise_output' % i, ep)

  def testCreationNoClasses(self):
    spec = copy.deepcopy(mobilenet_v2.V2_DEF)
    net, ep = mobilenet.mobilenet(
        tf.placeholder(tf.float32, (10, 224, 224, 16)), conv_defs=spec,
        num_classes=None)
    self.assertIs(net, ep['global_pool'])

  def testImageSizes(self):
    for input_size, output_size in [(224, 7), (192, 6), (160, 5),
                                    (128, 4), (96, 3)]:
      tf.reset_default_graph()
      _, ep = mobilenet_v2.mobilenet(
          tf.placeholder(tf.float32, (10, input_size, input_size, 3)))

      self.assertEqual(ep['layer_18/output'].get_shape().as_list()[1:3],
                       [output_size] * 2)

  def testWithSplits(self):
    spec = copy.deepcopy(mobilenet_v2.V2_DEF)
    spec['overrides'] = {
        (ops.expanded_conv,): dict(split_expansion=2),
    }
    _, _ = mobilenet.mobilenet(
        tf.placeholder(tf.float32, (10, 224, 224, 16)), conv_defs=spec)
    num_convs = len(find_ops('Conv2D'))
    # All but 3 op has 3 conv operatore, the remainign 3 have one
    # and there is one unaccounted.
    self.assertEqual(num_convs, len(spec['spec']) * 3 - 5)

  def testWithOutputStride8(self):
    out, _ = mobilenet.mobilenet_base(
        tf.placeholder(tf.float32, (10, 224, 224, 16)),
        conv_defs=mobilenet_v2.V2_DEF,
        output_stride=8,
        scope='MobilenetV2')
    self.assertEqual(out.get_shape().as_list()[1:3], [28, 28])

  def testDivisibleBy(self):
    tf.reset_default_graph()
    mobilenet_v2.mobilenet(
        tf.placeholder(tf.float32, (10, 224, 224, 16)),
        conv_defs=mobilenet_v2.V2_DEF,
        divisible_by=16,
        min_depth=32)
    s = [op.outputs[0].get_shape().as_list()[-1] for op in find_ops('Conv2D')]
    s = set(s)
    self.assertSameElements([32, 64, 96, 160, 192, 320, 384, 576, 960, 1280,
                             1001], s)

  def testDivisibleByWithArgScope(self):
    tf.reset_default_graph()
    # Verifies that depth_multiplier arg scope actually works
    # if no default min_depth is provided.
    with slim.arg_scope((mobilenet.depth_multiplier,), min_depth=32):
      mobilenet_v2.mobilenet(
          tf.placeholder(tf.float32, (10, 224, 224, 2)),
          conv_defs=mobilenet_v2.V2_DEF, depth_multiplier=0.1)
      s = [op.outputs[0].get_shape().as_list()[-1] for op in find_ops('Conv2D')]
      s = set(s)
      self.assertSameElements(s, [32, 192, 128, 1001])

  def testFineGrained(self):
    tf.reset_default_graph()
    # Verifies that depth_multiplier arg scope actually works
    # if no default min_depth is provided.

    mobilenet_v2.mobilenet(
        tf.placeholder(tf.float32, (10, 224, 224, 2)),
        conv_defs=mobilenet_v2.V2_DEF, depth_multiplier=0.01,
        finegrain_classification_mode=True)
    s = [op.outputs[0].get_shape().as_list()[-1] for op in find_ops('Conv2D')]
    s = set(s)
    # All convolutions will be 8->48, except for the last one.
    self.assertSameElements(s, [8, 48, 1001, 1280])

  def testMobilenetBase(self):
    tf.reset_default_graph()
    # Verifies that mobilenet_base returns pre-pooling layer.
    with slim.arg_scope((mobilenet.depth_multiplier,), min_depth=32):
      net, _ = mobilenet_v2.mobilenet_base(
          tf.placeholder(tf.float32, (10, 224, 224, 16)),
          conv_defs=mobilenet_v2.V2_DEF, depth_multiplier=0.1)
      self.assertEqual(net.get_shape().as_list(), [10, 7, 7, 128])

  def testWithOutputStride16(self):
    tf.reset_default_graph()
    out, _ = mobilenet.mobilenet_base(
        tf.placeholder(tf.float32, (10, 224, 224, 16)),
        conv_defs=mobilenet_v2.V2_DEF,
        output_stride=16)
    self.assertEqual(out.get_shape().as_list()[1:3], [14, 14])

  def testWithOutputStride8AndExplicitPadding(self):
    tf.reset_default_graph()
    out, _ = mobilenet.mobilenet_base(
        tf.placeholder(tf.float32, (10, 224, 224, 16)),
        conv_defs=mobilenet_v2.V2_DEF,
        output_stride=8,
        use_explicit_padding=True,
        scope='MobilenetV2')
    self.assertEqual(out.get_shape().as_list()[1:3], [28, 28])

  def testWithOutputStride16AndExplicitPadding(self):
    tf.reset_default_graph()
    out, _ = mobilenet.mobilenet_base(
        tf.placeholder(tf.float32, (10, 224, 224, 16)),
        conv_defs=mobilenet_v2.V2_DEF,
        output_stride=16,
        use_explicit_padding=True)
    self.assertEqual(out.get_shape().as_list()[1:3], [14, 14])

  def testBatchNormScopeDoesNotHaveIsTrainingWhenItsSetToNone(self):
    sc = mobilenet.training_scope(is_training=None)
    self.assertNotIn('is_training', sc[slim.arg_scope_func_key(
        slim.batch_norm)])

  def testBatchNormScopeDoesHasIsTrainingWhenItsNotNone(self):
    sc = mobilenet.training_scope(is_training=False)
    self.assertIn('is_training', sc[slim.arg_scope_func_key(slim.batch_norm)])
    sc = mobilenet.training_scope(is_training=True)
    self.assertIn('is_training', sc[slim.arg_scope_func_key(slim.batch_norm)])
    sc = mobilenet.training_scope()
    self.assertIn('is_training', sc[slim.arg_scope_func_key(slim.batch_norm)])


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

[((38, 7, 38, 29), 'tensorflow.get_default_graph', 'tf.get_default_graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((189, 2, 189, 16), 'tensorflow.test.main', 'tf.test.main', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((45, 4, 45, 28), 'tensorflow.reset_default_graph', 'tf.reset_default_graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((64, 11, 64, 45), 'copy.deepcopy', 'copy.deepcopy', ({(64, 25, 64, 44): 'mobilenet_v2.V2_DEF'}, {}), '(mobilenet_v2.V2_DEF)', False, 'import copy\n'), ((81, 11, 81, 45), 'copy.deepcopy', 'copy.deepcopy', ({(81, 25, 81, 44): 'mobilenet_v2.V2_DEF'}, {}), '(mobilenet_v2.V2_DEF)', False, 'import copy\n'), ((101, 4, 101, 28), 'tensorflow.reset_default_graph', 'tf.reset_default_graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((113, 4, 113, 28), 'tensorflow.reset_default_graph', 'tf.reset_default_graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((125, 4, 125, 28), 'tensorflow.reset_default_graph', 'tf.reset_default_graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((139, 4, 139, 28), 'tensorflow.reset_default_graph', 'tf.reset_default_graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((148, 4, 148, 28), 'tensorflow.reset_default_graph', 'tf.reset_default_graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((156, 4, 156, 28), 'tensorflow.reset_default_graph', 'tf.reset_default_graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((166, 4, 166, 28), 'tensorflow.reset_default_graph', 'tf.reset_default_graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((175, 9, 175, 51), 'nets.mobilenet.mobilenet.training_scope', 'mobilenet.training_scope', (), '', False, 'from nets.mobilenet import mobilenet\n'), ((180, 9, 180, 52), 'nets.mobilenet.mobilenet.training_scope', 'mobilenet.training_scope', (), '', False, 'from nets.mobilenet import mobilenet\n'), ((182, 9, 182, 51), 'nets.mobilenet.mobilenet.training_scope', 'mobilenet.training_scope', (), '', False, 'from nets.mobilenet import mobilenet\n'), ((184, 9, 184, 35), 'nets.mobilenet.mobilenet.training_scope', 'mobilenet.training_scope', ({}, {}), '()', False, 'from nets.mobilenet import mobilenet\n'), ((50, 8, 50, 54), 'tensorflow.placeholder', 'tf.placeholder', ({(50, 23, 50, 33): 'tf.float32', (50, 35, 50, 53): '(10, 224, 224, 16)'}, {}), '(tf.float32, (10, 224, 224, 16))', True, 'import tensorflow as tf\n'), ((66, 8, 66, 54), 'tensorflow.placeholder', 'tf.placeholder', ({(66, 23, 66, 33): 'tf.float32', (66, 35, 66, 53): '(10, 224, 224, 16)'}, {}), '(tf.float32, (10, 224, 224, 16))', True, 'import tensorflow as tf\n'), ((73, 6, 73, 30), 'tensorflow.reset_default_graph', 'tf.reset_default_graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((86, 8, 86, 54), 'tensorflow.placeholder', 'tf.placeholder', ({(86, 23, 86, 33): 'tf.float32', (86, 35, 86, 53): '(10, 224, 224, 16)'}, {}), '(tf.float32, (10, 224, 224, 16))', True, 'import tensorflow as tf\n'), ((94, 8, 94, 54), 'tensorflow.placeholder', 'tf.placeholder', ({(94, 23, 94, 33): 'tf.float32', (94, 35, 94, 53): '(10, 224, 224, 16)'}, {}), '(tf.float32, (10, 224, 224, 16))', True, 'import tensorflow as tf\n'), ((103, 8, 103, 54), 'tensorflow.placeholder', 'tf.placeholder', ({(103, 23, 103, 33): 'tf.float32', (103, 35, 103, 53): '(10, 224, 224, 16)'}, {}), '(tf.float32, (10, 224, 224, 16))', True, 'import tensorflow as tf\n'), ((130, 8, 130, 53), 'tensorflow.placeholder', 'tf.placeholder', ({(130, 23, 130, 33): 'tf.float32', (130, 35, 130, 52): '(10, 224, 224, 2)'}, {}), '(tf.float32, (10, 224, 224, 2))', True, 'import tensorflow as tf\n'), ((150, 8, 150, 54), 'tensorflow.placeholder', 'tf.placeholder', ({(150, 23, 150, 33): 'tf.float32', (150, 35, 150, 53): '(10, 224, 224, 16)'}, {}), '(tf.float32, (10, 224, 224, 16))', True, 'import tensorflow as tf\n'), ((158, 8, 158, 54), 'tensorflow.placeholder', 'tf.placeholder', ({(158, 23, 158, 33): 'tf.float32', (158, 35, 158, 53): '(10, 224, 224, 16)'}, {}), '(tf.float32, (10, 224, 224, 16))', True, 'import tensorflow as tf\n'), ((168, 8, 168, 54), 'tensorflow.placeholder', 'tf.placeholder', ({(168, 23, 168, 33): 'tf.float32', (168, 35, 168, 53): '(10, 224, 224, 16)'}, {}), '(tf.float32, (10, 224, 224, 16))', True, 'import tensorflow as tf\n'), ((75, 10, 75, 69), 'tensorflow.placeholder', 'tf.placeholder', ({(75, 25, 75, 35): 'tf.float32', (75, 37, 75, 68): '(10, input_size, input_size, 3)'}, {}), '(tf.float32, (10, input_size, input_size, 3))', True, 'import tensorflow as tf\n'), ((118, 10, 118, 55), 'tensorflow.placeholder', 'tf.placeholder', ({(118, 25, 118, 35): 'tf.float32', (118, 37, 118, 54): '(10, 224, 224, 2)'}, {}), '(tf.float32, (10, 224, 224, 2))', True, 'import tensorflow as tf\n'), ((143, 10, 143, 56), 'tensorflow.placeholder', 'tf.placeholder', ({(143, 25, 143, 35): 'tf.float32', (143, 37, 143, 55): '(10, 224, 224, 16)'}, {}), '(tf.float32, (10, 224, 224, 16))', True, 'import tensorflow as tf\n')]