Internal change

PiperOrigin-RevId: 351701695

official/vision/beta/serving/export_saved_model_lib.py

@@ -23,6 +23,7 @@ from official.core import train_utils
 from official.vision.beta import configs
 from official.vision.beta.serving import detection
 from official.vision.beta.serving import image_classification
+from official.vision.beta.serving import semantic_segmentation
 
 
 def export_inference_graph(input_type, batch_size, input_image_size, params,
@@ -62,15 +63,15 @@ def export_inference_graph(input_type, batch_size, input_image_size, params,
   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)
+        params=params, batch_size=batch_size, input_image_size=input_image_size)
   elif isinstance(params.task, configs.retinanet.RetinaNetTask) or isinstance(
       params.task, configs.maskrcnn.MaskRCNNTask):
     export_module = detection.DetectionModule(
-        params=params,
-        batch_size=batch_size,
-        input_image_size=input_image_size)
+        params=params, batch_size=batch_size, input_image_size=input_image_size)
+  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)
   else:
     raise ValueError('Export module not implemented for {} task.'.format(
         type(params.task)))

official/vision/beta/serving/semantic_segmentation.py

+# 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.
+# ==============================================================================
+"""Semantic segmentation input and model functions for serving/inference."""
+
+import tensorflow as tf
+
+from official.vision.beta.modeling import factory
+from official.vision.beta.ops import preprocess_ops
+from official.vision.beta.serving import export_base
+
+
+MEAN_RGB = (0.485 * 255, 0.456 * 255, 0.406 * 255)
+STDDEV_RGB = (0.229 * 255, 0.224 * 255, 0.225 * 255)
+
+
+class SegmentationModule(export_base.ExportModule):
+  """Segmentation Module."""
+
+  def build_model(self):
+    input_specs = tf.keras.layers.InputSpec(
+        shape=[self._batch_size] + self._input_image_size + [3])
+
+    self._model = factory.build_segmentation_model(
+        input_specs=input_specs,
+        model_config=self._params.task.model,
+        l2_regularizer=None)
+
+    return self._model
+
+  def _build_inputs(self, image):
+    """Builds classification model inputs for serving."""
+
+    # Normalizes image with mean and std pixel values.
+    image = preprocess_ops.normalize_image(image,
+                                           offset=MEAN_RGB,
+                                           scale=STDDEV_RGB)
+
+    image, _ = preprocess_ops.resize_and_crop_image(
+        image,
+        self._input_image_size,
+        padded_size=self._input_image_size,
+        aug_scale_min=1.0,
+        aug_scale_max=1.0)
+    return image
+
+  def _run_inference_on_image_tensors(self, images):
+    """Cast image to float and run inference.
+
+    Args:
+      images: uint8 Tensor of shape [batch_size, None, None, 3]
+    Returns:
+      Tensor holding classification output logits.
+    """
+    with tf.device('cpu:0'):
+      images = tf.cast(images, dtype=tf.float32)
+
+      images = tf.nest.map_structure(
+          tf.identity,
+          tf.map_fn(
+              self._build_inputs, elems=images,
+              fn_output_signature=tf.TensorSpec(
+                  shape=self._input_image_size + [3], dtype=tf.float32),
+              parallel_iterations=32
+              )
+          )
+
+    masks = self._model(images, training=False)
+    masks = tf.image.resize(masks, self._input_image_size, method='bilinear')
+
+    return dict(predicted_masks=masks)

official/vision/beta/serving/semantic_segmentation_test.py

+# 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.
+# ==============================================================================
+"""Test for semantic segmentation export lib."""
+
+import io
+import os
+
+from absl.testing import parameterized
+import numpy as np
+from PIL import Image
+import tensorflow as tf
+
+from official.common import registry_imports  # pylint: disable=unused-import
+from official.core import exp_factory
+from official.vision.beta.serving import semantic_segmentation
+
+
+class SemanticSegmentationExportTest(tf.test.TestCase, parameterized.TestCase):
+
+  def _get_segmentation_module(self):
+    params = exp_factory.get_exp_config('seg_deeplabv3_pascal')
+    params.task.model.backbone.dilated_resnet.model_id = 50
+    segmentation_module = semantic_segmentation.SegmentationModule(
+        params, batch_size=1, input_image_size=[112, 112])
+    return segmentation_module
+
+  def _export_from_module(self, module, input_type, save_directory):
+    if input_type == 'image_tensor':
+      input_signature = tf.TensorSpec(shape=[None, 112, 112, 3], dtype=tf.uint8)
+      signatures = {
+          'serving_default':
+              module.inference_from_image_tensors.get_concrete_function(
+                  input_signature)
+      }
+    elif input_type == 'image_bytes':
+      input_signature = tf.TensorSpec(shape=[None], dtype=tf.string)
+      signatures = {
+          'serving_default':
+              module.inference_from_image_bytes.get_concrete_function(
+                  input_signature)
+      }
+    elif input_type == 'tf_example':
+      input_signature = tf.TensorSpec(shape=[None], dtype=tf.string)
+      signatures = {
+          'serving_default':
+              module.inference_from_tf_example.get_concrete_function(
+                  input_signature)
+      }
+    else:
+      raise ValueError('Unrecognized `input_type`')
+
+    tf.saved_model.save(module,
+                        save_directory,
+                        signatures=signatures)
+
+  def _get_dummy_input(self, input_type):
+    """Get dummy input for the given input type."""
+
+    if input_type == 'image_tensor':
+      return tf.zeros((1, 112, 112, 3), dtype=np.uint8)
+    elif input_type == 'image_bytes':
+      image = Image.fromarray(np.zeros((112, 112, 3), dtype=np.uint8))
+      byte_io = io.BytesIO()
+      image.save(byte_io, 'PNG')
+      return [byte_io.getvalue()]
+    elif input_type == 'tf_example':
+      image_tensor = tf.zeros((112, 112, 3), dtype=tf.uint8)
+      encoded_jpeg = tf.image.encode_jpeg(tf.constant(image_tensor)).numpy()
+      example = tf.train.Example(
+          features=tf.train.Features(
+              feature={
+                  'image/encoded':
+                      tf.train.Feature(
+                          bytes_list=tf.train.BytesList(value=[encoded_jpeg])),
+              })).SerializeToString()
+      return [example]
+
+  @parameterized.parameters(
+      {'input_type': 'image_tensor'},
+      {'input_type': 'image_bytes'},
+      {'input_type': 'tf_example'},
+  )
+  def test_export(self, input_type='image_tensor'):
+    tmp_dir = self.get_temp_dir()
+
+    module = self._get_segmentation_module()
+    model = module.build_model()
+
+    self._export_from_module(module, input_type, tmp_dir)
+
+    self.assertTrue(os.path.exists(os.path.join(tmp_dir, 'saved_model.pb')))
+    self.assertTrue(os.path.exists(
+        os.path.join(tmp_dir, 'variables', 'variables.index')))
+    self.assertTrue(os.path.exists(
+        os.path.join(tmp_dir, 'variables', 'variables.data-00000-of-00001')))
+
+    imported = tf.saved_model.load(tmp_dir)
+    segmentation_fn = imported.signatures['serving_default']
+
+    images = self._get_dummy_input(input_type)
+    processed_images = tf.nest.map_structure(
+        tf.stop_gradient,
+        tf.map_fn(
+            module._build_inputs,
+            elems=tf.zeros((1, 112, 112, 3), dtype=tf.uint8),
+            fn_output_signature=tf.TensorSpec(
+                shape=[112, 112, 3], dtype=tf.float32)))
+    expected_output = tf.image.resize(
+        model(processed_images, training=False), [112, 112], method='bilinear')
+    out = segmentation_fn(tf.constant(images))
+    self.assertAllClose(out['predicted_masks'].numpy(), expected_output.numpy())
+
+if __name__ == '__main__':
+  tf.test.main()