No public description

PiperOrigin-RevId: 639040818

Author: tensorflower-gardener

official/vision/serving/detection.py

@@ -41,27 +41,37 @@ def _padded_size(self):
       return self._input_image_size
 
   def _build_model(self):
-
     nms_versions_supporting_dynamic_batch_size = {'batched', 'v2', 'v3'}
     nms_version = self.params.task.model.detection_generator.nms_version
-    if (self._batch_size is None and
-        nms_version not in nms_versions_supporting_dynamic_batch_size):
-      logging.info('nms_version is set to `batched` because `%s` '
-                   'does not support with dynamic batch size.', nms_version)
+    if (
+        self._batch_size is None
+        and nms_version not in nms_versions_supporting_dynamic_batch_size
+    ):
+      logging.info(
+          'nms_version is set to `batched` because `%s` '
+          'does not support with dynamic batch size.',
+          nms_version,
+      )
       self.params.task.model.detection_generator.nms_version = 'batched'
 
-    input_specs = tf_keras.layers.InputSpec(shape=[
-        self._batch_size, *self._padded_size, 3])
+    input_specs = tf_keras.layers.InputSpec(
+        shape=[self._batch_size, *self._padded_size, 3]
+    )
 
     if isinstance(self.params.task.model, configs.maskrcnn.MaskRCNN):
       model = factory.build_maskrcnn(
-          input_specs=input_specs, model_config=self.params.task.model)
+          input_specs=input_specs, model_config=self.params.task.model
+      )
     elif isinstance(self.params.task.model, configs.retinanet.RetinaNet):
       model = factory.build_retinanet(
-          input_specs=input_specs, model_config=self.params.task.model)
+          input_specs=input_specs, model_config=self.params.task.model
+      )
     else:
-      raise ValueError('Detection module not implemented for {} model.'.format(
-          type(self.params.task.model)))
+      raise ValueError(
+          'Detection module not implemented for {} model.'.format(
+              type(self.params.task.model)
+          )
+      )
 
     return model
 
@@ -73,7 +83,8 @@ def _build_anchor_boxes(self):
         max_level=model_params.max_level,
         num_scales=model_params.anchor.num_scales,
         aspect_ratios=model_params.anchor.aspect_ratios,
-        anchor_size=model_params.anchor.anchor_size)
+        anchor_size=model_params.anchor.anchor_size,
+    )
     return input_anchor(image_size=self._padded_size)
 
   def _build_inputs(self, image):
@@ -85,7 +96,8 @@ def _build_inputs(self, image):
 
     # Normalizes image with mean and std pixel values.
     image = preprocess_ops.normalize_image(
-        image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB)
+        image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB
+    )
 
     image, image_info = preprocess_ops.resize_and_crop_image(
         image,
@@ -131,20 +143,24 @@ def preprocess(
 
     Args:
       images: The images tensor.
+
     Returns:
       images: The images tensor cast to float.
       anchor_boxes: Dict mapping anchor levels to anchor boxes.
       image_info: Tensor containing the details of the image resizing.
-
     """
     model_params = self.params.task.model
     with tf.device('cpu:0'):
       # Tensor Specs for map_fn outputs (images, anchor_boxes, and image_info).
-      images_spec = tf.TensorSpec(shape=self._padded_size + [3],
-                                  dtype=tf.float32)
+      images_spec = tf.TensorSpec(
+          shape=self._padded_size + [3], dtype=tf.float32
+      )
 
-      num_anchors = model_params.anchor.num_scales * len(
-          model_params.anchor.aspect_ratios) * 4
+      num_anchors = (
+          model_params.anchor.num_scales
+          * len(model_params.anchor.aspect_ratios)
+          * 4
+      )
       anchor_shapes = []
       for level in range(model_params.min_level, model_params.max_level + 1):
         anchor_level_spec = tf.TensorSpec(
@@ -153,7 +169,8 @@ def preprocess(
                 math.ceil(self._padded_size[1] / 2**level),
                 num_anchors,
             ],
-            dtype=tf.float32)
+            dtype=tf.float32,
+        )
         anchor_shapes.append((str(level), anchor_level_spec))
 
       image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)
@@ -163,9 +180,14 @@ def preprocess(
           tf.map_fn(
               self._build_inputs,
               elems=images,
-              fn_output_signature=(images_spec, dict(anchor_shapes),
-                                   image_info_spec),
-              parallel_iterations=32))
+              fn_output_signature=(
+                  images_spec,
+                  dict(anchor_shapes),
+                  image_info_spec,
+              ),
+              parallel_iterations=32,
+          ),
+      )
 
       return images, anchor_boxes, image_info
 
@@ -174,6 +196,7 @@ def serve(self, images: tf.Tensor):
 
     Args:
       images: uint8 Tensor of shape [batch_size, None, None, 3]
+
     Returns:
       Tensor holding detection output logits.
     """
@@ -190,10 +213,15 @@ def serve(self, images: tf.Tensor):
         # [desired_height, desired_width], [y_scale, x_scale],
         # [y_offset, x_offset]]. When input_type is tflite, input image is
         # supposed to be preprocessed already.
-        image_info = tf.convert_to_tensor([[
-            self._input_image_size, self._input_image_size, [1.0, 1.0], [0, 0]
-        ]],
-                                          dtype=tf.float32)
+        image_info = tf.convert_to_tensor(
+            [[
+                self._input_image_size,
+                self._input_image_size,
+                [1.0, 1.0],
+                [0, 0],
+            ]],
+            dtype=tf.float32,
+        )
     input_image_shape = image_info[:, 1, :]
 
     # To overcome keras.Model extra limitation to save a model with layers that
@@ -226,20 +254,23 @@ def serve(self, images: tf.Tensor):
         # point outputs.
         if export_config.cast_num_detections_to_float:
           detections['num_detections'] = tf.cast(
-              detections['num_detections'], dtype=tf.float32)
+              detections['num_detections'], dtype=tf.float32
+          )
         if export_config.cast_detection_classes_to_float:
           detections['detection_classes'] = tf.cast(
-              detections['detection_classes'], dtype=tf.float32)
+              detections['detection_classes'], dtype=tf.float32
+          )
 
       final_outputs = {
           'detection_boxes': detections['detection_boxes'],
           'detection_scores': detections['detection_scores'],
           'detection_classes': detections['detection_classes'],
-          'num_detections': detections['num_detections']
+          'num_detections': detections['num_detections'],
       }
       if 'detection_outer_boxes' in detections:
-        final_outputs['detection_outer_boxes'] = (
-            detections['detection_outer_boxes'])
+        final_outputs['detection_outer_boxes'] = detections[
+            'detection_outer_boxes'
+        ]
     else:
       # For RetinaNet model, apply export_config.
       if isinstance(self.params.task.model, configs.retinanet.RetinaNet):
@@ -250,7 +281,7 @@ def serve(self, images: tf.Tensor):
           detections = self._normalize_coordinates(detections, keys, image_info)
       final_outputs = {
           'decoded_boxes': detections['decoded_boxes'],
-          'decoded_box_scores': detections['decoded_box_scores']
+          'decoded_box_scores': detections['decoded_box_scores'],
       }
 
     if 'detection_masks' in detections.keys():

official/vision/serving/export_saved_model_lib.py

@@ -15,7 +15,7 @@
 r"""Vision models export utility function for serving/inference."""
 
 import os
-from typing import Optional, List, Union, Text, Dict
+from typing import Dict, List, Optional, Union
 
 from absl import logging
 import tensorflow as tf, tf_keras
@@ -45,7 +45,7 @@ def export_inference_graph(
     log_model_flops_and_params: bool = False,
     checkpoint: Optional[tf.train.Checkpoint] = None,
     input_name: Optional[str] = None,
-    function_keys: Optional[Union[List[Text], Dict[Text, Text]]] = None,
+    function_keys: Optional[Union[List[str], Dict[str, str]]] = None,
     add_tpu_function_alias: Optional[bool] = False,
 ):
   """Exports inference graph for the model specified in the exp config.
@@ -83,57 +83,68 @@ def export_inference_graph(
 
   if export_checkpoint_subdir:
     output_checkpoint_directory = os.path.join(
-        export_dir, export_checkpoint_subdir)
+        export_dir, export_checkpoint_subdir
+    )
   else:
     output_checkpoint_directory = None
 
   if export_saved_model_subdir:
     output_saved_model_directory = os.path.join(
-        export_dir, export_saved_model_subdir)
+        export_dir, export_saved_model_subdir
+    )
   else:
     output_saved_model_directory = export_dir
 
   # TODO(arashwan): Offers a direct path to use ExportModule with Task objects.
   if not export_module:
-    if isinstance(params.task,
-                  configs.image_classification.ImageClassificationTask):
+    if isinstance(
+        params.task, configs.image_classification.ImageClassificationTask
+    ):
       export_module = image_classification.ClassificationModule(
           params=params,
           batch_size=batch_size,
           input_image_size=input_image_size,
           input_type=input_type,
           num_channels=num_channels,
-          input_name=input_name)
+          input_name=input_name,
+      )
     elif isinstance(params.task, configs.retinanet.RetinaNetTask) or isinstance(
-        params.task, configs.maskrcnn.MaskRCNNTask):
+        params.task, configs.maskrcnn.MaskRCNNTask
+    ):
       export_module = detection.DetectionModule(
           params=params,
           batch_size=batch_size,
           input_image_size=input_image_size,
           input_type=input_type,
           num_channels=num_channels,
-          input_name=input_name)
-    elif isinstance(params.task,
-                    configs.semantic_segmentation.SemanticSegmentationTask):
+          input_name=input_name,
+      )
+    elif isinstance(
+        params.task, configs.semantic_segmentation.SemanticSegmentationTask
+    ):
       export_module = semantic_segmentation.SegmentationModule(
           params=params,
           batch_size=batch_size,
           input_image_size=input_image_size,
           input_type=input_type,
           num_channels=num_channels,
-          input_name=input_name)
-    elif isinstance(params.task,
-                    configs.video_classification.VideoClassificationTask):
+          input_name=input_name,
+      )
+    elif isinstance(
+        params.task, configs.video_classification.VideoClassificationTask
+    ):
       export_module = video_classification.VideoClassificationModule(
           params=params,
           batch_size=batch_size,
           input_image_size=input_image_size,
           input_type=input_type,
           num_channels=num_channels,
-          input_name=input_name)
+          input_name=input_name,
+      )
     else:
-      raise ValueError('Export module not implemented for {} task.'.format(
-          type(params.task)))
+      raise ValueError(
+          'Export module not implemented for {} task.'.format(type(params.task))
+      )
 
   if add_tpu_function_alias:
     if input_type == 'image_tensor':
@@ -160,7 +171,8 @@ def export_inference_graph(
       checkpoint=checkpoint,
       checkpoint_path=checkpoint_path,
       timestamped=False,
-      save_options=save_options)
+      save_options=save_options,
+  )
 
   if output_checkpoint_directory:
     ckpt = tf.train.Checkpoint(model=export_module.model)
@@ -171,16 +183,16 @@ def export_inference_graph(
     inputs_kwargs = None
     if isinstance(
         params.task,
-        (configs.retinanet.RetinaNetTask, configs.maskrcnn.MaskRCNNTask)):
+        (configs.retinanet.RetinaNetTask, configs.maskrcnn.MaskRCNNTask),
+    ):
       # We need to create inputs_kwargs argument to specify the input shapes for
       # subclass model that overrides model.call to take multiple inputs,
       # e.g., RetinaNet model.
       inputs_kwargs = {
-          'images':
-              tf.TensorSpec([1] + input_image_size + [num_channels],
-                            tf.float32),
-          'image_shape':
-              tf.TensorSpec([1, 2], tf.float32)
+          'images': tf.TensorSpec(
+              [1] + input_image_size + [num_channels], tf.float32
+          ),
+          'image_shape': tf.TensorSpec([1, 2], tf.float32),
       }
       dummy_inputs = {
           k: tf.ones(v.shape.as_list(), tf.float32)
@@ -191,9 +203,14 @@ def export_inference_graph(
     else:
       logging.info(
           'Logging model flops and params not implemented for %s task.',
-          type(params.task))
+          type(params.task),
+      )
       return
-    train_utils.try_count_flops(export_module.model, inputs_kwargs,
-                                os.path.join(export_dir, 'model_flops.txt'))
-    train_utils.write_model_params(export_module.model,
-                                   os.path.join(export_dir, 'model_params.txt'))
+    train_utils.try_count_flops(
+        export_module.model,
+        inputs_kwargs,
+        os.path.join(export_dir, 'model_flops.txt'),
+    )
+    train_utils.write_model_params(
+        export_module.model, os.path.join(export_dir, 'model_params.txt')
+    )