Upload 4 files

export_tflite_graph_tf2.py

@@ -0,0 +1,152 @@
+# 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.
+# ==============================================================================
+r"""Exports TF2 detection SavedModel for conversion to TensorFlow Lite.
+Link to the TF2 Detection Zoo:
+https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md
+The output folder will contain an intermediate SavedModel that can be used with
+the TfLite converter.
+NOTE: This only supports SSD meta-architectures for now.
+One input:
+  image: a float32 tensor of shape[1, height, width, 3] containing the
+  *normalized* input image.
+  NOTE: See the `preprocess` function defined in the feature extractor class
+  in the object_detection/models directory.
+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
+Example Usage:
+--------------
+python object_detection/export_tflite_graph_tf2.py \
+    --pipeline_config_path path/to/ssd_model/pipeline.config \
+    --trained_checkpoint_dir path/to/ssd_model/checkpoint \
+    --output_directory path/to/exported_model_directory
+The expected output SavedModel would be in the directory
+path/to/exported_model_directory (which is created if it does not exist).
+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 1 (in which we change the NMS iou_threshold to be 0.5 and
+NMS score_threshold to be 0.0):
+python object_detection/export_tflite_model_tf2.py \
+    --pipeline_config_path path/to/ssd_model/pipeline.config \
+    --trained_checkpoint_dir path/to/ssd_model/checkpoint \
+    --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 \
+                } \
+             } \
+          } \
+       } \
+       "
+Example Usage 2 (export CenterNet model for keypoint estimation task with fixed
+shape resizer and customized input resolution):
+python object_detection/export_tflite_model_tf2.py \
+    --pipeline_config_path path/to/ssd_model/pipeline.config \
+    --trained_checkpoint_dir path/to/ssd_model/checkpoint \
+    --output_directory path/to/exported_model_directory \
+    --keypoint_label_map_path path/to/label_map.txt \
+    --max_detections 10 \
+    --centernet_include_keypoints true \
+    --config_override " \
+            model{ \
+              center_net { \
+                image_resizer { \
+                  fixed_shape_resizer { \
+                    height: 320 \
+                    width: 320 \
+                  } \
+                } \
+              } \
+            }" \
+"""
+from absl import app
+from absl import flags
+
+import tensorflow.compat.v2 as tf
+from google.protobuf import text_format
+from object_detection import export_tflite_graph_lib_tf2
+from object_detection.protos import pipeline_pb2
+
+tf.enable_v2_behavior()
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string(
+    'pipeline_config_path', None,
+    'Path to a pipeline_pb2.TrainEvalPipelineConfig config '
+    'file.')
+flags.DEFINE_string('trained_checkpoint_dir', None,
+                    'Path to trained checkpoint directory')
+flags.DEFINE_string('output_directory', None, 'Path to write outputs.')
+flags.DEFINE_string(
+    'config_override', '', 'pipeline_pb2.TrainEvalPipelineConfig '
+    'text proto to override pipeline_config_path.')
+flags.DEFINE_integer('max_detections', 10,
+                     'Maximum number of detections (boxes) to return.')
+# SSD-specific flags
+flags.DEFINE_bool(
+    'ssd_use_regular_nms', False,
+    'Flag to set postprocessing op to use Regular NMS instead of Fast NMS '
+    '(Default false).')
+# CenterNet-specific flags
+flags.DEFINE_bool(
+    'centernet_include_keypoints', False,
+    'Whether to export the predicted keypoint tensors. Only CenterNet model'
+    ' supports this flag.'
+)
+flags.DEFINE_string(
+    'keypoint_label_map_path', None,
+    'Path of the label map used by CenterNet keypoint estimation task. If'
+    ' provided, the label map path in the pipeline config will be replaced by'
+    ' this one. Note that it is only used when exporting CenterNet model for'
+    ' keypoint estimation task.'
+)
+
+
+def main(argv):
+  del argv  # Unused.
+  flags.mark_flag_as_required('pipeline_config_path')
+  flags.mark_flag_as_required('trained_checkpoint_dir')
+  flags.mark_flag_as_required('output_directory')
+
+  pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
+
+  with tf.io.gfile.GFile(FLAGS.pipeline_config_path, 'r') as f:
+    text_format.Parse(f.read(), pipeline_config)
+  override_config = pipeline_pb2.TrainEvalPipelineConfig()
+  text_format.Parse(FLAGS.config_override, override_config)
+  pipeline_config.MergeFrom(override_config)
+
+  export_tflite_graph_lib_tf2.export_tflite_model(
+      pipeline_config, FLAGS.trained_checkpoint_dir, FLAGS.output_directory,
+      FLAGS.max_detections, FLAGS.ssd_use_regular_nms,
+      FLAGS.centernet_include_keypoints, FLAGS.keypoint_label_map_path)
+
+
+if __name__ == '__main__':
+  app.run(main)

exporter_main_v2.py

@@ -0,0 +1,159 @@
+# 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.
+# ==============================================================================
+
+r"""Tool to export an object detection model for inference.
+
+Prepares an object detection tensorflow graph for inference using model
+configuration and a trained checkpoint. Outputs associated checkpoint files,
+a SavedModel, and a copy of the model config.
+
+The inference graph contains one of three input nodes depending on the user
+specified option.
+  * `image_tensor`: Accepts a uint8 4-D tensor of shape [1, None, None, 3]
+  * `float_image_tensor`: Accepts a float32 4-D tensor of shape
+    [1, None, None, 3]
+  * `encoded_image_string_tensor`: Accepts a 1-D string tensor of shape [None]
+    containing encoded PNG or JPEG images. Image resolutions are expected to be
+    the same if more than 1 image is provided.
+  * `tf_example`: Accepts a 1-D string tensor of shape [None] containing
+    serialized TFExample protos. Image resolutions are expected to be the same
+    if more than 1 image is provided.
+
+and the following output nodes returned by the model.postprocess(..):
+  * `num_detections`: Outputs float32 tensors of the form [batch]
+      that specifies the number of valid boxes per image in the batch.
+  * `detection_boxes`: Outputs float32 tensors of the form
+      [batch, num_boxes, 4] containing detected boxes.
+  * `detection_scores`: Outputs float32 tensors of the form
+      [batch, num_boxes] containing class scores for the detections.
+  * `detection_classes`: Outputs float32 tensors of the form
+      [batch, num_boxes] containing classes for the detections.
+
+
+Example Usage:
+--------------
+python exporter_main_v2.py \
+    --input_type image_tensor \
+    --pipeline_config_path path/to/ssd_inception_v2.config \
+    --trained_checkpoint_dir path/to/checkpoint \
+    --output_directory path/to/exported_model_directory
+    --use_side_inputs True/False \
+    --side_input_shapes dim_0,dim_1,...dim_a/.../dim_0,dim_1,...,dim_z \
+    --side_input_names name_a,name_b,...,name_c \
+    --side_input_types type_1,type_2
+
+The expected output would be in the directory
+path/to/exported_model_directory (which is created if it does not exist)
+holding two subdirectories (corresponding to checkpoint and SavedModel,
+respectively) and a copy of the pipeline config.
+
+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 second stage post-processing score
+threshold to be 0.5):
+
+python exporter_main_v2.py \
+    --input_type image_tensor \
+    --pipeline_config_path path/to/ssd_inception_v2.config \
+    --trained_checkpoint_dir path/to/checkpoint \
+    --output_directory path/to/exported_model_directory \
+    --config_override " \
+            model{ \
+              faster_rcnn { \
+                second_stage_post_processing { \
+                  batch_non_max_suppression { \
+                    score_threshold: 0.5 \
+                  } \
+                } \
+              } \
+            }"
+
+If side inputs are desired, the following arguments could be appended
+(the example below is for Context R-CNN).
+   --use_side_inputs True \
+   --side_input_shapes 1,2000,2057/1 \
+   --side_input_names context_features,valid_context_size \
+   --side_input_types tf.float32,tf.int32
+"""
+from absl import app
+from absl import flags
+
+import tensorflow.compat.v2 as tf
+from google.protobuf import text_format
+from object_detection import exporter_lib_v2
+from object_detection.protos import pipeline_pb2
+
+tf.enable_v2_behavior()
+
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string('input_type', 'image_tensor', 'Type of input node. Can be '
+                    'one of [`image_tensor`, `encoded_image_string_tensor`, '
+                    '`tf_example`, `float_image_tensor`]')
+flags.DEFINE_string('pipeline_config_path', None,
+                    'Path to a pipeline_pb2.TrainEvalPipelineConfig config '
+                    'file.')
+flags.DEFINE_string('trained_checkpoint_dir', None,
+                    'Path to trained checkpoint directory')
+flags.DEFINE_string('output_directory', None, 'Path to write outputs.')
+flags.DEFINE_string('config_override', '',
+                    'pipeline_pb2.TrainEvalPipelineConfig '
+                    'text proto to override pipeline_config_path.')
+flags.DEFINE_boolean('use_side_inputs', False,
+                     'If True, uses side inputs as well as image inputs.')
+flags.DEFINE_string('side_input_shapes', '',
+                    'If use_side_inputs is True, this explicitly sets '
+                    'the shape of the side input tensors to a fixed size. The '
+                    'dimensions are to be provided as a comma-separated list '
+                    'of integers. A value of -1 can be used for unknown '
+                    'dimensions. A `/` denotes a break, starting the shape of '
+                    'the next side input tensor. This flag is required if '
+                    'using side inputs.')
+flags.DEFINE_string('side_input_types', '',
+                    'If use_side_inputs is True, this explicitly sets '
+                    'the type of the side input tensors. The '
+                    'dimensions are to be provided as a comma-separated list '
+                    'of types, each of `string`, `integer`, or `float`. '
+                    'This flag is required if using side inputs.')
+flags.DEFINE_string('side_input_names', '',
+                    'If use_side_inputs is True, this explicitly sets '
+                    'the names of the side input tensors required by the model '
+                    'assuming the names will be a comma-separated list of '
+                    'strings. This flag is required if using side inputs.')
+
+flags.mark_flag_as_required('pipeline_config_path')
+flags.mark_flag_as_required('trained_checkpoint_dir')
+flags.mark_flag_as_required('output_directory')
+
+
+def main(_):
+  pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
+  with tf.io.gfile.GFile(FLAGS.pipeline_config_path, 'r') as f:
+    text_format.Merge(f.read(), pipeline_config)
+  text_format.Merge(FLAGS.config_override, pipeline_config)
+  exporter_lib_v2.export_inference_graph(
+      FLAGS.input_type, pipeline_config, FLAGS.trained_checkpoint_dir,
+      FLAGS.output_directory, FLAGS.use_side_inputs, FLAGS.side_input_shapes,
+      FLAGS.side_input_types, FLAGS.side_input_names)
+
+
+if __name__ == '__main__':
+  app.run(main)

model_main_tf2.py

@@ -0,0 +1,113 @@
+# 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.
+# ==============================================================================
+
+r"""Creates and runs TF2 object detection models.
+
+For local training/evaluation run:
+PIPELINE_CONFIG_PATH=path/to/pipeline.config
+MODEL_DIR=/tmp/model_outputs
+NUM_TRAIN_STEPS=10000
+SAMPLE_1_OF_N_EVAL_EXAMPLES=1
+python model_main_tf2.py -- \
+  --model_dir=$MODEL_DIR --num_train_steps=$NUM_TRAIN_STEPS \
+  --sample_1_of_n_eval_examples=$SAMPLE_1_OF_N_EVAL_EXAMPLES \
+  --pipeline_config_path=$PIPELINE_CONFIG_PATH \
+  --alsologtostderr
+"""
+from absl import flags
+import tensorflow.compat.v2 as tf
+from object_detection import model_lib_v2
+
+flags.DEFINE_string('pipeline_config_path', None, 'Path to pipeline config '
+                    'file.')
+flags.DEFINE_integer('num_train_steps', None, 'Number of train steps.')
+flags.DEFINE_bool('eval_on_train_data', False, 'Enable evaluating on train '
+                  'data (only supported in distributed training).')
+flags.DEFINE_integer('sample_1_of_n_eval_examples', None, 'Will sample one of '
+                     'every n eval input examples, where n is provided.')
+flags.DEFINE_integer('sample_1_of_n_eval_on_train_examples', 5, 'Will sample '
+                     'one of every n train input examples for evaluation, '
+                     'where n is provided. This is only used if '
+                     '`eval_training_data` is True.')
+flags.DEFINE_string(
+    'model_dir', None, 'Path to output model directory '
+                       'where event and checkpoint files will be written.')
+flags.DEFINE_string(
+    'checkpoint_dir', None, 'Path to directory holding a checkpoint.  If '
+    '`checkpoint_dir` is provided, this binary operates in eval-only mode, '
+    'writing resulting metrics to `model_dir`.')
+
+flags.DEFINE_integer('eval_timeout', 3600, 'Number of seconds to wait for an'
+                     'evaluation checkpoint before exiting.')
+
+flags.DEFINE_bool('use_tpu', False, 'Whether the job is executing on a TPU.')
+flags.DEFINE_string(
+    'tpu_name',
+    default=None,
+    help='Name of the Cloud TPU for Cluster Resolvers.')
+flags.DEFINE_integer(
+    'num_workers', 1, 'When num_workers > 1, training uses '
+    'MultiWorkerMirroredStrategy. When num_workers = 1 it uses '
+    'MirroredStrategy.')
+flags.DEFINE_integer(
+    'checkpoint_every_n', 1000, 'Integer defining how often we checkpoint.')
+flags.DEFINE_boolean('record_summaries', True,
+                     ('Whether or not to record summaries during'
+                      ' training.'))
+
+FLAGS = flags.FLAGS
+
+
+def main(unused_argv):
+  flags.mark_flag_as_required('model_dir')
+  flags.mark_flag_as_required('pipeline_config_path')
+  tf.config.set_soft_device_placement(True)
+
+  if FLAGS.checkpoint_dir:
+    model_lib_v2.eval_continuously(
+        pipeline_config_path=FLAGS.pipeline_config_path,
+        model_dir=FLAGS.model_dir,
+        train_steps=FLAGS.num_train_steps,
+        sample_1_of_n_eval_examples=FLAGS.sample_1_of_n_eval_examples,
+        sample_1_of_n_eval_on_train_examples=(
+            FLAGS.sample_1_of_n_eval_on_train_examples),
+        checkpoint_dir=FLAGS.checkpoint_dir,
+        wait_interval=300, timeout=FLAGS.eval_timeout)
+  else:
+    if FLAGS.use_tpu:
+      # TPU is automatically inferred if tpu_name is None and
+      # we are running under cloud ai-platform.
+      resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
+          FLAGS.tpu_name)
+      tf.config.experimental_connect_to_cluster(resolver)
+      tf.tpu.experimental.initialize_tpu_system(resolver)
+      strategy = tf.distribute.experimental.TPUStrategy(resolver)
+    elif FLAGS.num_workers > 1:
+      strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
+    else:
+      strategy = tf.compat.v2.distribute.MirroredStrategy()
+
+    with strategy.scope():
+      model_lib_v2.train_loop(
+          pipeline_config_path=FLAGS.pipeline_config_path,
+          model_dir=FLAGS.model_dir,
+          train_steps=FLAGS.num_train_steps,
+          use_tpu=FLAGS.use_tpu,
+          checkpoint_every_n=FLAGS.checkpoint_every_n,
+          record_summaries=FLAGS.record_summaries)
+
+if __name__ == '__main__':
+  tf.compat.v1.app.run()