text: # 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.
  # ==============================================================================
  """Library to export TFLite-compatible SavedModel from TF2 detection models."""
  import os
  import numpy as np
  import tensorflow.compat.v1 as tf1
  import tensorflow.compat.v2 as tf
  
  from object_detection.builders import model_builder
  from object_detection.builders import post_processing_builder
  from object_detection.core import box_list
  from object_detection.core import standard_fields as fields
  
  _DEFAULT_NUM_CHANNELS = 3
  _DEFAULT_NUM_COORD_BOX = 4
  _MAX_CLASSES_PER_DETECTION = 1
  _DETECTION_POSTPROCESS_FUNC = 'TFLite_Detection_PostProcess'
  
  
  def get_const_center_size_encoded_anchors(anchors):
    """Exports center-size encoded anchors as a constant tensor.
  
    Args:
      anchors: a float32 tensor of shape [num_anchors, 4] containing the anchor
        boxes
  
    Returns:
      encoded_anchors: a float32 constant tensor of shape [num_anchors, 4]
      containing the anchor boxes.
    """
    anchor_boxlist = box_list.BoxList(anchors)
    y, x, h, w = anchor_boxlist.get_center_coordinates_and_sizes()
    num_anchors = y.get_shape().as_list()
  
    with tf1.Session() as sess:
      y_out, x_out, h_out, w_out = sess.run([y, x, h, w])
    encoded_anchors = tf1.constant(
        np.transpose(np.stack((y_out, x_out, h_out, w_out))),
        dtype=tf1.float32,
        shape=[num_anchors[0], _DEFAULT_NUM_COORD_BOX],
        name='anchors')
    return num_anchors[0], encoded_anchors
  
  
  class SSDModule(tf.Module):
    """Inference Module for TFLite-friendly SSD models."""
  
    def __init__(self, pipeline_config, detection_model, max_detections,
                 use_regular_nms):
      """Initialization.
  
      Args:
        pipeline_config: The original pipeline_pb2.TrainEvalPipelineConfig
        detection_model: The detection model to use for inference.
        max_detections: Max detections desired from the TFLite model.
        use_regular_nms: If True, TFLite model uses the (slower) multi-class NMS.
      """
      self._process_config(pipeline_config)
      self._pipeline_config = pipeline_config
      self._model = detection_model
      self._max_detections = max_detections
      self._use_regular_nms = use_regular_nms
  
    def _process_config(self, pipeline_config):
      self._num_classes = pipeline_config.model.ssd.num_classes
      self._nms_score_threshold = pipeline_config.model.ssd.post_processing.batch_non_max_suppression.score_threshold
      self._nms_iou_threshold = pipeline_config.model.ssd.post_processing.batch_non_max_suppression.iou_threshold
      self._scale_values = {}
      self._scale_values[
          'y_scale'] = pipeline_config.model.ssd.box_coder.faster_rcnn_box_coder.y_scale
      self._scale_values[
          'x_scale'] = pipeline_config.model.ssd.box_coder.faster_rcnn_box_coder.x_scale
      self._scale_values[
          'h_scale'] = pipeline_config.model.ssd.box_coder.faster_rcnn_box_coder.height_scale
      self._scale_values[
          'w_scale'] = pipeline_config.model.ssd.box_coder.faster_rcnn_box_coder.width_scale
  
      image_resizer_config = pipeline_config.model.ssd.image_resizer
      image_resizer = image_resizer_config.WhichOneof('image_resizer_oneof')
      self._num_channels = _DEFAULT_NUM_CHANNELS
  
      if image_resizer == 'fixed_shape_resizer':
        self._height = image_resizer_config.fixed_shape_resizer.height
        self._width = image_resizer_config.fixed_shape_resizer.width
        if image_resizer_config.fixed_shape_resizer.convert_to_grayscale:
          self._num_channels = 1
      else:
        raise ValueError(
            'Only fixed_shape_resizer'
            'is supported with tflite. Found {}'.format(
                image_resizer_config.WhichOneof('image_resizer_oneof')))
  
    def input_shape(self):
      """Returns shape of TFLite model input."""
      return [1, self._height, self._width, self._num_channels]
  
    def postprocess_implements_signature(self):
      """Returns tf.implements signature for MLIR legalization of TFLite NMS."""
      implements_signature = [
          'name: "%s"' % _DETECTION_POSTPROCESS_FUNC,
          'attr { key: "max_detections" value { i: %d } }' % self._max_detections,
          'attr { key: "max_classes_per_detection" value { i: %d } }' %
          _MAX_CLASSES_PER_DETECTION,
          'attr { key: "use_regular_nms" value { b: %s } }' %
          str(self._use_regular_nms).lower(),
          'attr { key: "nms_score_threshold" value { f: %f } }' %
          self._nms_score_threshold,
          'attr { key: "nms_iou_threshold" value { f: %f } }' %
          self._nms_iou_threshold,
          'attr { key: "y_scale" value { f: %f } }' %
          self._scale_values['y_scale'],
          'attr { key: "x_scale" value { f: %f } }' %
          self._scale_values['x_scale'],
          'attr { key: "h_scale" value { f: %f } }' %
          self._scale_values['h_scale'],
          'attr { key: "w_scale" value { f: %f } }' %
          self._scale_values['w_scale'],
          'attr { key: "num_classes" value { i: %d } }' % self._num_classes
      ]
      implements_signature = ' '.join(implements_signature)
      return implements_signature
  
    def _get_postprocess_fn(self, num_anchors, num_classes):
      # There is no TF equivalent for TFLite's custom post-processing op.
      # So we add an 'empty' composite function here, that is legalized to the
      # custom op with MLIR.
      @tf.function(
          experimental_implements=self.postprocess_implements_signature())
      # pylint: disable=g-unused-argument,unused-argument
      def dummy_post_processing(box_encodings, class_predictions, anchors):
        boxes = tf.constant(0.0, dtype=tf.float32, name='boxes')
        scores = tf.constant(0.0, dtype=tf.float32, name='scores')
        classes = tf.constant(0.0, dtype=tf.float32, name='classes')
        num_detections = tf.constant(0.0, dtype=tf.float32, name='num_detections')
        return boxes, scores, classes, num_detections
  
      return dummy_post_processing
  
    @tf.function
    def inference_fn(self, image):
      """Encapsulates SSD inference for TFLite conversion.
  
      NOTE: The Args & Returns sections below indicate the TFLite model signature,
      and not what the TF graph does (since the latter does not include the custom
      NMS op used by TFLite)
  
      Args:
        image: a float32 tensor of shape [num_anchors, 4] containing the anchor
          boxes
  
      Returns:
        num_detections: a float32 scalar denoting number of total detections.
        classes: a float32 tensor denoting class ID for each detection.
        scores: a float32 tensor denoting score for each detection.
        boxes: a float32 tensor denoting coordinates of each detected box.
      """
      predicted_tensors = self._model.predict(image, true_image_shapes=None)
      # The score conversion occurs before the post-processing custom op
      _, score_conversion_fn = post_processing_builder.build(
          self._pipeline_config.model.ssd.post_processing)
      class_predictions = score_conversion_fn(
          predicted_tensors['class_predictions_with_background'])
  
      with tf.name_scope('raw_outputs'):
        # 'raw_outputs/box_encodings': a float32 tensor of shape
        #   [1, num_anchors, 4] containing the encoded box predictions. Note that
        #   these are raw predictions and no Non-Max suppression is applied on
        #   them and no decode center size boxes is applied to them.
        box_encodings = tf.identity(
            predicted_tensors['box_encodings'], name='box_encodings')
        # 'raw_outputs/class_predictions': a float32 tensor of shape
        #   [1, num_anchors, num_classes] containing the class scores for each
        #   anchor after applying score conversion.
        class_predictions = tf.identity(
            class_predictions, name='class_predictions')
      # 'anchors': a float32 tensor of shape
      #   [4, num_anchors] containing the anchors as a constant node.
      num_anchors, anchors = get_const_center_size_encoded_anchors(
          predicted_tensors['anchors'])
      anchors = tf.identity(anchors, name='anchors')
  
      # tf.function@ seems to reverse order of inputs, so reverse them here.
      return self._get_postprocess_fn(num_anchors,
                                      self._num_classes)(box_encodings,
                                                         class_predictions,
                                                         anchors)[::-1]
  
  
  class CenterNetModule(tf.Module):
    """Inference Module for TFLite-friendly CenterNet models.
  
    The exported CenterNet model includes the preprocessing and postprocessing
    logics so the caller should pass in the raw image pixel values. It supports
    both object detection and keypoint estimation task.
    """
  
    def __init__(self, pipeline_config, max_detections, include_keypoints):
      """Initialization.
  
      Args:
        pipeline_config: The original pipeline_pb2.TrainEvalPipelineConfig
        max_detections: Max detections desired from the TFLite model.
        include_keypoints: If set true, the output dictionary will include the
          keypoint coordinates and keypoint confidence scores.
      """
      self._max_detections = max_detections
      self._include_keypoints = include_keypoints
      self._process_config(pipeline_config)
      self._pipeline_config = pipeline_config
      self._model = model_builder.build(
          self._pipeline_config.model, is_training=False)
  
    def get_model(self):
      return self._model
  
    def _process_config(self, pipeline_config):
      self._num_classes = pipeline_config.model.center_net.num_classes
  
      center_net_config = pipeline_config.model.center_net
      image_resizer_config = center_net_config.image_resizer
      image_resizer = image_resizer_config.WhichOneof('image_resizer_oneof')
      self._num_channels = _DEFAULT_NUM_CHANNELS
  
      if image_resizer == 'fixed_shape_resizer':
        self._height = image_resizer_config.fixed_shape_resizer.height
        self._width = image_resizer_config.fixed_shape_resizer.width
        if image_resizer_config.fixed_shape_resizer.convert_to_grayscale:
          self._num_channels = 1
      else:
        raise ValueError(
            'Only fixed_shape_resizer'
            'is supported with tflite. Found {}'.format(image_resizer))
  
      center_net_config.object_center_params.max_box_predictions = (
          self._max_detections)
  
      if not self._include_keypoints:
        del center_net_config.keypoint_estimation_task[:]
  
    def input_shape(self):
      """Returns shape of TFLite model input."""
      return [1, self._height, self._width, self._num_channels]
  
    @tf.function
    def inference_fn(self, image):
      """Encapsulates CenterNet inference for TFLite conversion.
  
      Args:
        image: a float32 tensor of shape [1, image_height, image_width, channel]
          denoting the image pixel values.
  
      Returns:
        A dictionary of predicted tensors:
          classes: a float32 tensor with shape [1, max_detections] denoting class
            ID for each detection.
          scores: a float32 tensor with shape [1, max_detections] denoting score
            for each detection.
          boxes: a float32 tensor with shape [1, max_detections, 4] denoting
            coordinates of each detected box.
          keypoints: a float32 with shape [1, max_detections, num_keypoints, 2]
            denoting the predicted keypoint coordinates (normalized in between
            0-1). Note that [:, :, :, 0] represents the y coordinates and
            [:, :, :, 1] represents the x coordinates.
          keypoint_scores: a float32 with shape [1, max_detections, num_keypoints]
            denoting keypoint confidence scores.
      """
      image = tf.cast(image, tf.float32)
      image, shapes = self._model.preprocess(image)
      prediction_dict = self._model.predict(image, None)
      detections = self._model.postprocess(
          prediction_dict, true_image_shapes=shapes)
  
      field_names = fields.DetectionResultFields
      classes_field = field_names.detection_classes
      classes = tf.cast(detections[classes_field], tf.float32)
      num_detections = tf.cast(detections[field_names.num_detections], tf.float32)
  
      if self._include_keypoints:
        model_outputs = (detections[field_names.detection_boxes], classes,
                         detections[field_names.detection_scores], num_detections,
                         detections[field_names.detection_keypoints],
                         detections[field_names.detection_keypoint_scores])
      else:
        model_outputs = (detections[field_names.detection_boxes], classes,
                         detections[field_names.detection_scores], num_detections)
  
      # tf.function@ seems to reverse order of inputs, so reverse them here.
      return model_outputs[::-1]
  
  
  def export_tflite_model(pipeline_config, trained_checkpoint_dir,
                          output_directory, max_detections, use_regular_nms,
                          include_keypoints=False):
    """Exports inference SavedModel for TFLite conversion.
  
    NOTE: Only supports SSD meta-architectures for now, and the output model will
    have static-shaped, single-batch input.
  
    This function creates `output_directory` if it does not already exist,
    which will hold the intermediate SavedModel that can be used with the TFLite
    converter.
  
    Args:
      pipeline_config: pipeline_pb2.TrainAndEvalPipelineConfig proto.
      trained_checkpoint_dir: Path to the trained checkpoint file.
      output_directory: Path to write outputs.
      max_detections: Max detections desired from the TFLite model.
      use_regular_nms: If True, TFLite model uses the (slower) multi-class NMS.
        Note that this argument is only used by the SSD model.
      include_keypoints: Decides whether to also output the keypoint predictions.
        Note that this argument is only used by the CenterNet model.
  
    Raises:
      ValueError: if pipeline is invalid.
    """
    output_saved_model_directory = os.path.join(output_directory, 'saved_model')
  
    # Build the underlying model using pipeline config.
    # TODO(b/162842801): Add support for other architectures.
    if pipeline_config.model.WhichOneof('model') == 'ssd':
      detection_model = model_builder.build(
          pipeline_config.model, is_training=False)
      ckpt = tf.train.Checkpoint(model=detection_model)
      # The module helps build a TF SavedModel appropriate for TFLite conversion.
      detection_module = SSDModule(pipeline_config, detection_model,
                                   max_detections, use_regular_nms)
    elif pipeline_config.model.WhichOneof('model') == 'center_net':
      detection_module = CenterNetModule(
          pipeline_config, max_detections, include_keypoints)
      ckpt = tf.train.Checkpoint(model=detection_module.get_model())
    else:
      raise ValueError('Only ssd or center_net models are supported in tflite. '
                       'Found {} in config'.format(
                           pipeline_config.model.WhichOneof('model')))
  
    manager = tf.train.CheckpointManager(
        ckpt, trained_checkpoint_dir, max_to_keep=1)
    status = ckpt.restore(manager.latest_checkpoint).expect_partial()
  
    # Getting the concrete function traces the graph and forces variables to
    # be constructed; only after this can we save the saved model.
    status.assert_existing_objects_matched()
    concrete_function = detection_module.inference_fn.get_concrete_function(
        tf.TensorSpec(
            shape=detection_module.input_shape(), dtype=tf.float32, name='input'))
    status.assert_existing_objects_matched()
  
    # Export SavedModel.
    tf.saved_model.save(
        detection_module,
        output_saved_model_directory,
        signatures=concrete_function)
  