text: # Copyright 2019 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.
  # ==============================================================================
  r"""Exports an LSTM detection model to use with tf-lite.
  
  Outputs file:
  * A tflite compatible frozen graph - $output_directory/tflite_graph.pb
  
  The exported graph has the following input and output nodes.
  
  Inputs:
  'input_video_tensor': a float32 tensor of shape
  [unroll_length, height, width, 3] containing the normalized input image.
  Note that the height and width must be compatible with the height and
  width configured in the fixed_shape_image resizer options in the pipeline
  config proto.
  
  Outputs:
  If add_postprocessing_op is true: frozen graph adds a
    TFLite_Detection_PostProcess custom op node has four outputs:
    detection_boxes: a float32 tensor of shape [1, num_boxes, 4] with box
    locations
    detection_classes: a float32 tensor of shape [1, num_boxes]
    with class indices
    detection_scores: a float32 tensor of shape [1, num_boxes]
    with class scores
    num_boxes: a float32 tensor of size 1 containing the number of detected boxes
  else:
    the graph has three outputs:
     'raw_outputs/box_encodings': a float32 tensor of shape [1, num_anchors, 4]
      containing the encoded box predictions.
     '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.
     'anchors': a float32 constant tensor of shape [num_anchors, 4]
      containing the anchor boxes.
  
  Example Usage:
  --------------
  python lstm_object_detection/export_tflite_lstd_graph.py \
      --pipeline_config_path path/to/lstm_pipeline.config \
      --trained_checkpoint_prefix path/to/model.ckpt \
      --output_directory path/to/exported_model_directory
  
  The expected output would be in the directory
  path/to/exported_model_directory (which is created if it does not exist)
  with contents:
   - tflite_graph.pbtxt
   - tflite_graph.pb
  Config overrides (see the `config_override` flag) are text protobufs
  (also of type pipeline_pb2.TrainEvalPipelineConfig) which are used to override
  certain fields in the provided pipeline_config_path.  These are useful for
  making small changes to the inference graph that differ from the training or
  eval config.
  
  Example Usage (in which we change the NMS iou_threshold to be 0.5 and
  NMS score_threshold to be 0.0):
  python lstm_object_detection/export_tflite_lstd_graph.py \
      --pipeline_config_path path/to/lstm_pipeline.config \
      --trained_checkpoint_prefix path/to/model.ckpt \
      --output_directory path/to/exported_model_directory
      --config_override " \
              model{ \
              ssd{ \
                post_processing { \
                  batch_non_max_suppression { \
                          score_threshold: 0.0 \
                          iou_threshold: 0.5 \
                  } \
               } \
            } \
         } \
         "
  """
  
  import tensorflow as tf
  
  from lstm_object_detection import export_tflite_lstd_graph_lib
  from lstm_object_detection.utils import config_util
  
  flags = tf.app.flags
  flags.DEFINE_string('output_directory', None, 'Path to write outputs.')
  flags.DEFINE_string(
      'pipeline_config_path', None,
      'Path to a pipeline_pb2.TrainEvalPipelineConfig config '
      'file.')
  flags.DEFINE_string('trained_checkpoint_prefix', None, 'Checkpoint prefix.')
  flags.DEFINE_integer('max_detections', 10,
                       'Maximum number of detections (boxes) to show.')
  flags.DEFINE_integer('max_classes_per_detection', 1,
                       'Maximum number of classes to output per detection box.')
  flags.DEFINE_integer(
      'detections_per_class', 100,
      'Number of anchors used per class in Regular Non-Max-Suppression.')
  flags.DEFINE_bool('add_postprocessing_op', True,
                    'Add TFLite custom op for postprocessing to the graph.')
  flags.DEFINE_bool(
      'use_regular_nms', False,
      'Flag to set postprocessing op to use Regular NMS instead of Fast NMS.')
  flags.DEFINE_string(
      'config_override', '', 'pipeline_pb2.TrainEvalPipelineConfig '
      'text proto to override pipeline_config_path.')
  
  FLAGS = flags.FLAGS
  
  
  def main(argv):
    del argv  # Unused.
    flags.mark_flag_as_required('output_directory')
    flags.mark_flag_as_required('pipeline_config_path')
    flags.mark_flag_as_required('trained_checkpoint_prefix')
  
    pipeline_config = config_util.get_configs_from_pipeline_file(
        FLAGS.pipeline_config_path)
  
    export_tflite_lstd_graph_lib.export_tflite_graph(
        pipeline_config,
        FLAGS.trained_checkpoint_prefix,
        FLAGS.output_directory,
        FLAGS.add_postprocessing_op,
        FLAGS.max_detections,
        FLAGS.max_classes_per_detection,
        use_regular_nms=FLAGS.use_regular_nms)
  
  
  if __name__ == '__main__':
    tf.app.run(main)
  