text: # Copyright 2020 Google Research. 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"""Inference related utilities."""
  import copy
  import os
  import time
  from typing import Text, Dict, Any
  from absl import logging
  import numpy as np
  import tensorflow as tf
  
  import dataloader
  import hparams_config
  import utils
  from keras import efficientdet_keras
  from keras import label_util
  from keras import util_keras
  from visualize import vis_utils
  
  
  def visualize_image(image,
                      boxes,
                      classes,
                      scores,
                      label_map=None,
                      min_score_thresh=0.01,
                      max_boxes_to_draw=1000,
                      line_thickness=2,
                      **kwargs):
    """Visualizes a given image.
  
    Args:
      image: a image with shape [H, W, C].
      boxes: a box prediction with shape [N, 4] ordered [ymin, xmin, ymax, xmax].
      classes: a class prediction with shape [N].
      scores: A list of float value with shape [N].
      label_map: a dictionary from class id to name.
      min_score_thresh: minimal score for showing. If claass probability is below
        this threshold, then the object will not show up.
      max_boxes_to_draw: maximum bounding box to draw.
      line_thickness: how thick is the bounding box line.
      **kwargs: extra parameters.
  
    Returns:
      output_image: an output image with annotated boxes and classes.
    """
    label_map = label_util.get_label_map(label_map or 'coco')
    category_index = {k: {'id': k, 'name': label_map[k]} for k in label_map}
    img = np.array(image)
    vis_utils.visualize_boxes_and_labels_on_image_array(
        img,
        boxes,
        classes,
        scores,
        category_index,
        min_score_thresh=min_score_thresh,
        max_boxes_to_draw=max_boxes_to_draw,
        line_thickness=line_thickness,
        **kwargs)
    return img
  
  
  class ExportNetwork(tf.Module):
  
    def __init__(self, model):
      super().__init__()
      self.model = model
  
    @tf.function
    def __call__(self, imgs):
      return tf.nest.flatten(self.model(imgs, training=False))
  
  
  class ExportModel(tf.Module):
    """Model to be exported as SavedModel/TFLite format."""
  
    def __init__(self, model, pre_mode='infer'):
      super().__init__()
      self.model = model
      self.pre_mode = pre_mode
  
    @tf.function
    def __call__(self, imgs):
      return self.model(
          imgs, training=False, pre_mode=self.pre_mode, post_mode='global')
  
  
  class ServingDriver:
    """A driver for serving single or batch images.
  
    This driver supports serving with image files or arrays, with configurable
    batch size.
  
    Example 1. Serving streaming image contents:
  
      driver = inference.ServingDriver(
        'efficientdet-d0', '/tmp/efficientdet-d0', batch_size=1)
      driver.build()
      for m in image_iterator():
        predictions = driver.serve_files([m])
        boxes, scores, classes, _ = tf.nest.map_structure(np.array, predictions)
        driver.visualize(m, boxes[0], scores[0], classes[0])
        # m is the new image with annotated boxes.
  
    Example 2. Serving batch image contents:
  
      imgs = []
      for f in ['/tmp/1.jpg', '/tmp/2.jpg']:
        imgs.append(np.array(Image.open(f)))
  
      driver = inference.ServingDriver(
        'efficientdet-d0', '/tmp/efficientdet-d0', batch_size=len(imgs))
      driver.build()
      predictions = driver.serve(imgs)
      boxes, scores, classes, _ = tf.nest.map_structure(np.array, predictions)
      for i in range(len(imgs)):
        driver.visualize(imgs[i], boxes[i], scores[i], classes[i])
  
    Example 3: another way is to use SavedModel:
  
      # step1: export a model.
      driver = inference.ServingDriver('efficientdet-d0', '/tmp/efficientdet-d0')
      driver.build()
      driver.export('/tmp/saved_model_path')
  
      # step2: Serve a model.
      driver.load(self.saved_model_dir)
      raw_images = []
      for f in tf.io.gfile.glob('/tmp/images/*.jpg'):
        raw_images.append(np.array(PIL.Image.open(f)))
      detections = driver.serve(raw_images)
      boxes, scores, classes, _ = tf.nest.map_structure(np.array, detections)
      for i in range(len(imgs)):
        driver.visualize(imgs[i], boxes[i], scores[i], classes[i])
    """
  
    def __init__(self,
                 model_name: Text,
                 ckpt_path: Text = None,
                 batch_size: int = 1,
                 only_network: bool = False,
                 model_params: Dict[Text, Any] = None):
      """Initialize the inference driver.
  
      Args:
        model_name: target model name, such as efficientdet-d0.
        ckpt_path: checkpoint path, such as /tmp/efficientdet-d0/.
        batch_size: batch size for inference.
        only_network: only use the network without pre/post processing.
        model_params: model parameters for overriding the config.
      """
      super().__init__()
      self.model_name = model_name
      self.ckpt_path = ckpt_path
      self.batch_size = batch_size
      self.only_network = only_network
  
      self.params = hparams_config.get_detection_config(model_name).as_dict()
  
      if model_params:
        self.params.update(model_params)
      self.params.update(dict(is_training_bn=False))
      self.label_map = self.params.get('label_map', None)
  
      self._model = None
  
      mixed_precision = self.params.get('mixed_precision', None)
      precision = utils.get_precision(
          self.params.get('strategy', None), mixed_precision)
      policy = tf.keras.mixed_precision.Policy(precision)
      tf.keras.mixed_precision.set_global_policy(policy)
  
    @property
    def model(self):
      if not self._model:
        self.build()
      return self._model
  
    @model.setter
    def model(self, model):
      self._model = model
  
    def build(self, params_override=None):
      """Build model and restore checkpoints."""
      params = copy.deepcopy(self.params)
      if params_override:
        params.update(params_override)
      config = hparams_config.get_efficientdet_config(self.model_name)
      config.override(params)
      if self.only_network:
        self.model = efficientdet_keras.EfficientDetNet(config=config)
      else:
        self.model = efficientdet_keras.EfficientDetModel(config=config)
      image_size = utils.parse_image_size(params['image_size'])
      self.model.build((self.batch_size, *image_size, 3))
      util_keras.restore_ckpt(self.model, self.ckpt_path,
                              self.params['moving_average_decay'],
                              skip_mismatch=False)
  
    def visualize(self, image, boxes, classes, scores, **kwargs):
      """Visualize prediction on image."""
      return visualize_image(image, boxes, classes.astype(int), scores,
                             self.label_map, **kwargs)
  
    def benchmark(self, image_arrays, bm_runs=10, trace_filename=None):
      """Benchmark inference latency/throughput.
  
      Args:
        image_arrays: a list of images in numpy array format.
        bm_runs: Number of benchmark runs.
        trace_filename: If None, specify the filename for saving trace.
      """
      _, spec = self._get_model_and_spec()
  
      @tf.function(input_signature=[spec])
      def test_func(image_arrays):
        return self.model(image_arrays)  # pylint: disable=not-callable
  
      for _ in range(3):  # warmup 3 runs.
        test_func(image_arrays)
  
      start = time.perf_counter()
      for _ in range(bm_runs):
        test_func(image_arrays)
      end = time.perf_counter()
      inference_time = (end - start) / bm_runs
  
      print('Per batch inference time: ', inference_time)
      print('FPS: ', self.batch_size / inference_time)
  
      if trace_filename:
        options = tf.profiler.experimental.ProfilerOptions()
        tf.profiler.experimental.start(trace_filename, options)
        test_func(image_arrays)
        tf.profiler.experimental.stop()
  
    def serve(self, image_arrays):
      """Serve a list of image arrays.
  
      Args:
        image_arrays: A list of image content with each image has shape [height,
          width, 3] and uint8 type.
  
      Returns:
        A list of detections.
      """
      if isinstance(self.model, tf.lite.Interpreter):
        input_details = self.model.get_input_details()
        output_details = self.model.get_output_details()
        self.model.set_tensor(input_details[0]['index'], np.array(image_arrays))
        self.model.invoke()
        return [self.model.get_tensor(x['index']) for x in output_details]
      return self.model(image_arrays)  # pylint: disable=not-callable
  
    def load(self, saved_model_dir_or_frozen_graph: Text):
      """Load the model using saved model or a frozen graph."""
      # Load saved model if it is a folder.
      if tf.saved_model.contains_saved_model(saved_model_dir_or_frozen_graph):
        self.model = tf.saved_model.load(saved_model_dir_or_frozen_graph)
        return
  
      if saved_model_dir_or_frozen_graph.endswith('.tflite'):
        self.model = tf.lite.Interpreter(saved_model_dir_or_frozen_graph)
        self.model.allocate_tensors()
        return
  
      # Load a frozen graph.
      def wrap_frozen_graph(graph_def, inputs, outputs):
        # https://www.tensorflow.org/guide/migrate
        imports_graph_def_fn = lambda: tf.import_graph_def(graph_def, name='')
        wrapped_import = tf.compat.v1.wrap_function(imports_graph_def_fn, [])
        import_graph = wrapped_import.graph
        return wrapped_import.prune(
            tf.nest.map_structure(import_graph.as_graph_element, inputs),
            tf.nest.map_structure(import_graph.as_graph_element, outputs))
  
      graph_def = tf.Graph().as_graph_def()
      with tf.io.gfile.GFile(saved_model_dir_or_frozen_graph, 'rb') as f:
        graph_def.ParseFromString(f.read())
  
      self.model = wrap_frozen_graph(
          graph_def,
          inputs='images:0',
          outputs=['Identity:0', 'Identity_1:0', 'Identity_2:0', 'Identity_3:0'])
  
    def freeze(self, func):
      """Freeze the graph."""
      # pylint: disable=g-import-not-at-top,disable=g-direct-tensorflow-import
      from tensorflow.python.framework.convert_to_constants \
        import convert_variables_to_constants_v2_as_graph
      _, graphdef = convert_variables_to_constants_v2_as_graph(func)
      return graphdef
  
    def _get_model_and_spec(self, tflite=None):
      """Get model instance and export spec."""
      if self.only_network or tflite:
        image_size = utils.parse_image_size(self.params['image_size'])
        spec = tf.TensorSpec(
            shape=[self.batch_size, *image_size, 3],
            dtype=tf.float32,
            name='images')
        if self.only_network:
          export_model = ExportNetwork(self.model)
        else:
          # If export tflite, we should remove preprocessing since TFLite doesn't
          # support dynamic shape.
          logging.info('Export model without preprocessing.')
          export_model = ExportModel(self.model, pre_mode=None)
        return export_model, spec
      else:
        spec = tf.TensorSpec(
            shape=[self.batch_size, None, None, 3], dtype=tf.uint8, name='images')
        export_model = ExportModel(self.model)
        return export_model, spec
  
    def export(self,
               output_dir: Text = None,
               tensorrt: Text = None,
               tflite: Text = None,
               file_pattern: Text = None,
               num_calibration_steps: int = 2000):
      """Export a saved model, frozen graph, and potential tflite/tensorrt model.
  
      Args:
        output_dir: the output folder for saved model.
        tensorrt: If not None, must be {'FP32', 'FP16', 'INT8'}.
        tflite: Type for post-training quantization.
        file_pattern: Glob for tfrecords, e.g. coco/val-*.tfrecord.
        num_calibration_steps: Number of post-training quantization calibration
          steps to run.
      """
      export_model, input_spec = self._get_model_and_spec(tflite)
      image_size = utils.parse_image_size(self.params['image_size'])
      if output_dir:
        tf.saved_model.save(
            export_model,
            output_dir,
            signatures=export_model.__call__.get_concrete_function(input_spec))
        logging.info('Model saved at %s', output_dir)
  
        # also save freeze pb file.
        graphdef = self.freeze(
            export_model.__call__.get_concrete_function(input_spec))
        proto_path = tf.io.write_graph(
            graphdef, output_dir, self.model_name + '_frozen.pb', as_text=False)
        logging.info('Frozen graph saved at %s', proto_path)
  
      if tflite:
        shape = (self.batch_size, *image_size, 3)
        input_spec = tf.TensorSpec(
            shape=shape, dtype=input_spec.dtype, name=input_spec.name)
        converter = tf.lite.TFLiteConverter.from_concrete_functions(
            [export_model.__call__.get_concrete_function(input_spec)])
        if tflite == 'FP32':
          converter.optimizations = [tf.lite.Optimize.DEFAULT]
          converter.target_spec.supported_types = [tf.float32]
        elif tflite == 'FP16':
          converter.optimizations = [tf.lite.Optimize.DEFAULT]
          converter.target_spec.supported_types = [tf.float16]
        elif tflite == 'INT8':
          if file_pattern:
            config = hparams_config.get_efficientdet_config(self.model_name)
            config.override(self.params)
            ds = dataloader.InputReader(
                file_pattern,
                is_training=False,
                max_instances_per_image=config.max_instances_per_image)(
                    config, batch_size=self.batch_size)
  
            def representative_dataset_gen():
              for image, _ in ds.take(num_calibration_steps):
                yield [image]
          else:  # Used for debugging, can remove later.
            logging.warn('Use real representative dataset instead of fake ones.')
            num_calibration_steps = 10
            def representative_dataset_gen():  # rewrite this for real data.
              for _ in range(num_calibration_steps):
                yield [tf.ones(shape, dtype=input_spec.dtype)]
  
          converter.representative_dataset = representative_dataset_gen
          converter.optimizations = [tf.lite.Optimize.DEFAULT]
          converter.inference_input_type = tf.uint8
          converter.inference_output_type = tf.uint8
          supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
          if not self.only_network:
            supported_ops.append(tf.lite.OpsSet.TFLITE_BUILTINS)
          converter.target_spec.supported_ops = supported_ops
  
        else:
          raise ValueError(f'Invalid tflite {tflite}: must be FP32, FP16, INT8.')
  
        tflite_path = os.path.join(output_dir, tflite.lower() + '.tflite')
        tflite_model = converter.convert()
        tf.io.gfile.GFile(tflite_path, 'wb').write(tflite_model)
        logging.info('TFLite is saved at %s', tflite_path)
  
      if tensorrt:
        trt_path = os.path.join(output_dir, 'tensorrt_' + tensorrt.lower())
        conversion_params = tf.experimental.tensorrt.ConversionParams(
            max_workspace_size_bytes=(2 << 20),
            maximum_cached_engines=1,
            precision_mode=tensorrt.upper())
        converter = tf.experimental.tensorrt.Converter(
            output_dir, conversion_params=conversion_params)
        converter.convert()
        converter.save(trt_path)
        logging.info('TensorRT model is saved at %s', trt_path)
  