Spaces:
Running
Running
# 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 | |
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 \ | |
} \ | |
} \ | |
} \ | |
}" | |
""" | |
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.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) | |
if __name__ == '__main__': | |
app.run(main) | |