# Lint as: python2, 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 model_builder under TensorFlow 2.X."""

import os
import unittest

import tensorflow.compat.v1 as tf

from google.protobuf import text_format
from object_detection.builders import model_builder
from object_detection.builders import model_builder_test
from object_detection.core import losses
from object_detection.models import center_net_resnet_feature_extractor
from object_detection.protos import center_net_pb2
from object_detection.protos import model_pb2
from object_detection.utils import tf_version


@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
class ModelBuilderTF2Test(model_builder_test.ModelBuilderTest):

  def default_ssd_feature_extractor(self):
    return 'ssd_resnet50_v1_fpn_keras'

  def default_faster_rcnn_feature_extractor(self):
    return 'faster_rcnn_resnet101_keras'

  def ssd_feature_extractors(self):
    return model_builder.SSD_KERAS_FEATURE_EXTRACTOR_CLASS_MAP

  def get_override_base_feature_extractor_hyperparams(self, extractor_type):
    return extractor_type in {}

  def faster_rcnn_feature_extractors(self):
    return model_builder.FASTER_RCNN_KERAS_FEATURE_EXTRACTOR_CLASS_MAP

  def get_fake_label_map_file_path(self):
    keypoint_spec_text = """
    item {
      name: "/m/01g317"
      id: 1
      display_name: "person"
      keypoints {
        id: 0
        label: 'nose'
      }
      keypoints {
        id: 1
        label: 'left_shoulder'
      }
      keypoints {
        id: 2
        label: 'right_shoulder'
      }
      keypoints {
        id: 3
        label: 'hip'
      }
    }
    """
    keypoint_label_map_path = os.path.join(
        self.get_temp_dir(), 'keypoint_label_map')
    with tf.gfile.Open(keypoint_label_map_path, 'wb') as f:
      f.write(keypoint_spec_text)
    return keypoint_label_map_path

  def get_fake_keypoint_proto(self):
    task_proto_txt = """
      task_name: "human_pose"
      task_loss_weight: 0.9
      keypoint_regression_loss_weight: 1.0
      keypoint_heatmap_loss_weight: 0.1
      keypoint_offset_loss_weight: 0.5
      heatmap_bias_init: 2.14
      keypoint_class_name: "/m/01g317"
      loss {
        classification_loss {
          penalty_reduced_logistic_focal_loss {
            alpha: 3.0
            beta: 4.0
          }
        }
        localization_loss {
          l1_localization_loss {
          }
        }
      }
      keypoint_label_to_std {
        key: "nose"
        value: 0.3
      }
      keypoint_label_to_std {
        key: "hip"
        value: 0.0
      }
      keypoint_candidate_score_threshold: 0.3
      num_candidates_per_keypoint: 12
      peak_max_pool_kernel_size: 5
      unmatched_keypoint_score: 0.05
      box_scale: 1.7
      candidate_search_scale: 0.2
      candidate_ranking_mode: "score_distance_ratio"
      offset_peak_radius: 3
      per_keypoint_offset: true
    """
    config = text_format.Merge(task_proto_txt,
                               center_net_pb2.CenterNet.KeypointEstimation())
    return config

  def get_fake_object_center_proto(self):
    proto_txt = """
      object_center_loss_weight: 0.5
      heatmap_bias_init: 3.14
      min_box_overlap_iou: 0.2
      max_box_predictions: 15
      classification_loss {
        penalty_reduced_logistic_focal_loss {
          alpha: 3.0
          beta: 4.0
        }
      }
    """
    return text_format.Merge(proto_txt,
                             center_net_pb2.CenterNet.ObjectCenterParams())

  def get_fake_object_detection_proto(self):
    proto_txt = """