Merge branch 'tensorflow:master' into exp_pr2

official/vision/beta/serving/export_base_v2.py

+# Copyright 2021 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.
+
+"""Base class for model export."""
+
+from typing import Dict, Optional, Text, Callable, Any, Union
+
+import tensorflow as tf
+
+from official.core import export_base
+
+
+class ExportModule(export_base.ExportModule):
+  """Base Export Module."""
+
+  def __init__(self,
+               params,
+               model: tf.keras.Model,
+               input_signature: Union[tf.TensorSpec, Dict[str, tf.TensorSpec]],
+               preprocessor: Optional[Callable[..., Any]] = None,
+               inference_step: Optional[Callable[..., Any]] = None,
+               postprocessor: Optional[Callable[..., Any]] = None):
+    """Initializes a module for export.
+
+    Args:
+      params: A dataclass for parameters to the module.
+      model: A tf.keras.Model instance to be exported.
+      input_signature: tf.TensorSpec, e.g.
+        tf.TensorSpec(shape=[None, 224, 224, 3], dtype=tf.uint8)
+      preprocessor: An optional callable function to preprocess the inputs.
+      inference_step: An optional callable function to forward-pass the model.
+      postprocessor: An optional callable function to postprocess the model
+        outputs.
+    """
+    super().__init__(params, model=model, inference_step=inference_step)
+    self.preprocessor = preprocessor
+    self.postprocessor = postprocessor
+    self.input_signature = input_signature
+
+  @tf.function
+  def serve(self, inputs):
+    x = self.preprocessor(inputs=inputs) if self.preprocessor else inputs
+    x = self.inference_step(x)
+    x = self.postprocessor(x) if self.postprocessor else x
+    return x
+
+  def get_inference_signatures(self, function_keys: Dict[Text, Text]):
+    """Gets defined function signatures.
+
+    Args:
+      function_keys: A dictionary with keys as the function to create signature
+        for and values as the signature keys when returns.
+
+    Returns:
+      A dictionary with key as signature key and value as concrete functions
+        that can be used for tf.saved_model.save.
+    """
+    signatures = {}
+    for _, def_name in function_keys.items():
+      signatures[def_name] = self.serve.get_concrete_function(
+          self.input_signature)
+    return signatures

official/vision/beta/serving/export_base_v2_test.py

+# Copyright 2021 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.
+
+"""Tests for official.core.export_base_v2."""
+import os
+
+import tensorflow as tf
+
+from official.core import export_base
+from official.vision.beta.serving import export_base_v2
+
+
+class TestModel(tf.keras.Model):
+
+  def __init__(self):
+    super().__init__()
+    self._dense = tf.keras.layers.Dense(2)
+
+  def call(self, inputs):
+    return {'outputs': self._dense(inputs)}
+
+
+class ExportBaseTest(tf.test.TestCase):
+
+  def test_preprocessor(self):
+    tmp_dir = self.get_temp_dir()
+    model = TestModel()
+    inputs = tf.ones([2, 4], tf.float32)
+
+    preprocess_fn = lambda inputs: 2 * inputs
+
+    module = export_base_v2.ExportModule(
+        params=None,
+        input_signature=tf.TensorSpec(shape=[2, 4]),
+        model=model,
+        preprocessor=preprocess_fn)
+    expected_output = model(preprocess_fn(inputs))
+    ckpt_path = tf.train.Checkpoint(model=model).save(
+        os.path.join(tmp_dir, 'ckpt'))
+    export_dir = export_base.export(
+        module, ['serving_default'],
+        export_savedmodel_dir=tmp_dir,
+        checkpoint_path=ckpt_path,
+        timestamped=False)
+    imported = tf.saved_model.load(export_dir)
+    output = imported.signatures['serving_default'](inputs)
+    print('output', output)
+    self.assertAllClose(
+        output['outputs'].numpy(), expected_output['outputs'].numpy())
+
+  def test_postprocessor(self):
+    tmp_dir = self.get_temp_dir()
+    model = TestModel()
+    inputs = tf.ones([2, 4], tf.float32)
+
+    postprocess_fn = lambda logits: {'outputs': 2 * logits['outputs']}
+
+    module = export_base_v2.ExportModule(
+        params=None,
+        model=model,
+        input_signature=tf.TensorSpec(shape=[2, 4]),
+        postprocessor=postprocess_fn)
+    expected_output = postprocess_fn(model(inputs))
+    ckpt_path = tf.train.Checkpoint(model=model).save(
+        os.path.join(tmp_dir, 'ckpt'))
+    export_dir = export_base.export(
+        module, ['serving_default'],
+        export_savedmodel_dir=tmp_dir,
+        checkpoint_path=ckpt_path,
+        timestamped=False)
+    imported = tf.saved_model.load(export_dir)
+    output = imported.signatures['serving_default'](inputs)
+    self.assertAllClose(
+        output['outputs'].numpy(), expected_output['outputs'].numpy())
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/vision/beta/serving/export_module_factory.py

+# Copyright 2021 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.
+
+"""Factory for vision export modules."""
+
+from typing import List, Optional
+
+import tensorflow as tf
+
+from official.core import config_definitions as cfg
+from official.vision.beta import configs
+from official.vision.beta.dataloaders import classification_input
+from official.vision.beta.modeling import factory
+from official.vision.beta.serving import export_base_v2 as export_base
+from official.vision.beta.serving import export_utils
+
+
+def create_classification_export_module(params: cfg.ExperimentConfig,
+                                        input_type: str,
+                                        batch_size: int,
+                                        input_image_size: List[int],
+                                        num_channels: int = 3):
+  """Creats classification export module."""
+  input_signature = export_utils.get_image_input_signatures(
+      input_type, batch_size, input_image_size, num_channels)
+  input_specs = tf.keras.layers.InputSpec(
+      shape=[batch_size] + input_image_size + [num_channels])
+
+  model = factory.build_classification_model(
+      input_specs=input_specs,
+      model_config=params.task.model,
+      l2_regularizer=None)
+
+  def preprocess_fn(inputs):
+    image_tensor = export_utils.parse_image(
+        inputs, input_type, input_image_size, num_channels)
+
+    def preprocess_image_fn(inputs):
+      return classification_input.Parser.inference_fn(
+          inputs, input_image_size, num_channels)
+
+    images = tf.map_fn(
+        preprocess_image_fn, elems=image_tensor,
+        fn_output_signature=tf.TensorSpec(
+            shape=input_image_size + [num_channels],
+            dtype=tf.float32))
+
+    return images
+
+  def postprocess_fn(logits):
+    probs = tf.nn.softmax(logits)
+    return {'logits': logits, 'probs': probs}
+
+  export_module = export_base.ExportModule(params,
+                                           model=model,
+                                           input_signature=input_signature,
+                                           preprocessor=preprocess_fn,
+                                           postprocessor=postprocess_fn)
+  return export_module
+
+
+def get_export_module(params: cfg.ExperimentConfig,
+                      input_type: str,
+                      batch_size: Optional[int],
+                      input_image_size: List[int],
+                      num_channels: int = 3) -> export_base.ExportModule:
+  """Factory for export modules."""
+  if isinstance(params.task,
+                configs.image_classification.ImageClassificationTask):
+    export_module = create_classification_export_module(
+        params, input_type, batch_size, input_image_size, num_channels)
+  else:
+    raise ValueError('Export module not implemented for {} task.'.format(
+        type(params.task)))
+  return export_module

official/vision/beta/serving/export_module_factory_test.py

+# Copyright 2021 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 vision modules."""
+
+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.core import export_base
+from official.vision.beta.dataloaders import classification_input
+from official.vision.beta.serving import export_module_factory
+
+
+class ImageClassificationExportTest(tf.test.TestCase, parameterized.TestCase):
+
+  def _get_classification_module(self, input_type, input_image_size):
+    params = exp_factory.get_exp_config('resnet_imagenet')
+    params.task.model.backbone.resnet.model_id = 18
+    module = export_module_factory.create_classification_export_module(
+        params, input_type, batch_size=1, input_image_size=input_image_size)
+    return module
+
+  def _get_dummy_input(self, input_type):
+    """Get dummy input for the given input type."""
+
+    if input_type == 'image_tensor':
+      return tf.zeros((1, 32, 32, 3), dtype=np.uint8)
+    elif input_type == 'image_bytes':
+      image = Image.fromarray(np.zeros((32, 32, 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((32, 32, 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'):
+    input_image_size = [32, 32]
+    tmp_dir = self.get_temp_dir()
+    module = self._get_classification_module(input_type, input_image_size)
+    # Test that the model restores any attrs that are trackable objects
+    # (eg: tables, resource variables, keras models/layers, tf.hub modules).
+    module.model.test_trackable = tf.keras.layers.InputLayer(input_shape=(4,))
+    ckpt_path = tf.train.Checkpoint(model=module.model).save(
+        os.path.join(tmp_dir, 'ckpt'))
+    export_dir = export_base.export(
+        module, [input_type],
+        export_savedmodel_dir=tmp_dir,
+        checkpoint_path=ckpt_path,
+        timestamped=False)
+
+    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(export_dir)
+    classification_fn = imported.signatures['serving_default']
+
+    images = self._get_dummy_input(input_type)
+
+    def preprocess_image_fn(inputs):
+      return classification_input.Parser.inference_fn(
+          inputs, input_image_size, num_channels=3)
+
+    processed_images = tf.map_fn(
+        preprocess_image_fn,
+        elems=tf.zeros([1] + input_image_size + [3], dtype=tf.uint8),
+        fn_output_signature=tf.TensorSpec(
+            shape=input_image_size + [3], dtype=tf.float32))
+    expected_logits = module.model(processed_images, training=False)
+    expected_prob = tf.nn.softmax(expected_logits)
+    out = classification_fn(tf.constant(images))
+
+    # The imported model should contain any trackable attrs that the original
+    # model had.
+    self.assertTrue(hasattr(imported.model, 'test_trackable'))
+    self.assertAllClose(
+        out['logits'].numpy(), expected_logits.numpy(), rtol=1e-04, atol=1e-04)
+    self.assertAllClose(
+        out['probs'].numpy(), expected_prob.numpy(), rtol=1e-04, atol=1e-04)
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/vision/beta/serving/export_saved_model_lib_v2.py

+# Copyright 2021 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"""Vision models export utility function for serving/inference."""
+
+import os
+from typing import Optional, List
+
+import tensorflow as tf
+
+from official.core import config_definitions as cfg
+from official.core import export_base
+from official.core import train_utils
+from official.vision.beta.serving import export_module_factory
+
+
+def export(
+    input_type: str,
+    batch_size: Optional[int],
+    input_image_size: List[int],
+    params: cfg.ExperimentConfig,
+    checkpoint_path: str,
+    export_dir: str,
+    num_channels: Optional[int] = 3,
+    export_module: Optional[export_base.ExportModule] = None,
+    export_checkpoint_subdir: Optional[str] = None,
+    export_saved_model_subdir: Optional[str] = None,
+    save_options: Optional[tf.saved_model.SaveOptions] = None):
+  """Exports the model specified in the exp config.
+
+  Saved model is stored at export_dir/saved_model, checkpoint is saved
+  at export_dir/checkpoint, and params is saved at export_dir/params.yaml.
+
+  Args:
+    input_type: One of `image_tensor`, `image_bytes`, `tf_example`.
+    batch_size: 'int', or None.
+    input_image_size: List or Tuple of height and width.
+    params: Experiment params.
+    checkpoint_path: Trained checkpoint path or directory.
+    export_dir: Export directory path.
+    num_channels: The number of input image channels.
+    export_module: Optional export module to be used instead of using params
+      to create one. If None, the params will be used to create an export
+      module.
+    export_checkpoint_subdir: Optional subdirectory under export_dir
+      to store checkpoint.
+    export_saved_model_subdir: Optional subdirectory under export_dir
+      to store saved model.
+    save_options: `SaveOptions` for `tf.saved_model.save`.
+  """
+
+  if export_checkpoint_subdir:
+    output_checkpoint_directory = os.path.join(
+        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)
+  else:
+    output_saved_model_directory = export_dir
+
+  export_module = export_module_factory.get_export_module(
+      params,
+      input_type=input_type,
+      batch_size=batch_size,
+      input_image_size=input_image_size,
+      num_channels=num_channels)
+
+  export_base.export(
+      export_module,
+      function_keys=[input_type],
+      export_savedmodel_dir=output_saved_model_directory,
+      checkpoint_path=checkpoint_path,
+      timestamped=False,
+      save_options=save_options)
+
+  if output_checkpoint_directory:
+    ckpt = tf.train.Checkpoint(model=export_module.model)
+    ckpt.save(os.path.join(output_checkpoint_directory, 'ckpt'))
+  train_utils.serialize_config(params, export_dir)

official/vision/beta/serving/export_utils.py

+# Copyright 2021 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.
+
+"""Helper utils for export library."""
+
+from typing import List, Optional
+import tensorflow as tf
+
+# pylint: disable=g-long-lambda
+
+
+def get_image_input_signatures(input_type: str,
+                               batch_size: Optional[int],
+                               input_image_size: List[int],
+                               num_channels: int = 3):
+  """Gets input signatures for an image.
+
+  Args:
+    input_type: A `str`, can be either tf_example, image_bytes, or image_tensor.
+    batch_size: `int` for batch size or None.
+    input_image_size: List[int] for the height and width of the input image.
+    num_channels: `int` for number of channels in the input image.
+  Returns:
+    tf.TensorSpec of the input tensor.
+  """
+  if input_type == 'image_tensor':
+    input_signature = tf.TensorSpec(
+        shape=[batch_size] + [None] * len(input_image_size) + [num_channels],
+        dtype=tf.uint8)
+  elif input_type in ['image_bytes', 'serve_examples', 'tf_example']:
+    input_signature = tf.TensorSpec(shape=[batch_size], dtype=tf.string)
+  elif input_type == 'tflite':
+    input_signature = tf.TensorSpec(
+        shape=[1] + input_image_size + [num_channels], dtype=tf.float32)
+  else:
+    raise ValueError('Unrecognized `input_type`')
+  return input_signature
+
+
+def decode_image(encoded_image_bytes: str,
+                 input_image_size: List[int],
+                 num_channels: int = 3,) -> tf.Tensor:
+  """Decodes an image bytes to an image tensor.
+
+  Use `tf.image.decode_image` to decode an image if input is expected to be 2D
+  image; otherwise use `tf.io.decode_raw` to convert the raw bytes to tensor
+  and reshape it to desire shape.
+
+  Args:
+    encoded_image_bytes: An encoded image string to be decoded.
+    input_image_size: List[int] for the desired input size. This will be used to
+      infer whether the image is 2d or 3d.
+    num_channels: `int` for number of image channels.
+
+  Returns:
+    A decoded image tensor.
+  """
+  if len(input_image_size) == 2:
+    # Decode an image if 2D input is expected.
+    image_tensor = tf.image.decode_image(
+        encoded_image_bytes, channels=num_channels)
+  else:
+    # Convert raw bytes into a tensor and reshape it, if not 2D input.
+    image_tensor = tf.io.decode_raw(encoded_image_bytes, out_type=tf.uint8)
+  image_tensor.set_shape([None] * len(input_image_size) + [num_channels])
+  return image_tensor
+
+
+def decode_image_tf_example(tf_example_string_tensor: tf.train.Example,
+                            input_image_size: List[int],
+                            num_channels: int = 3,
+                            encoded_key: str = 'image/encoded'
+                            ) -> tf.Tensor:
+  """Decodes a TF Example to an image tensor."""
+
+  keys_to_features = {
+      encoded_key: tf.io.FixedLenFeature((), tf.string, default_value=''),
+  }
+  parsed_tensors = tf.io.parse_single_example(
+      serialized=tf_example_string_tensor, features=keys_to_features)
+  image_tensor = decode_image(
+      parsed_tensors[encoded_key],
+      input_image_size=input_image_size,
+      num_channels=num_channels)
+  return image_tensor
+
+
+def parse_image(
+    inputs, input_type: str, input_image_size: List[int], num_channels: int):
+  """Parses image."""
+  if input_type in ['tf_example', 'serve_examples']:
+    decode_image_tf_example_fn = (
+        lambda x: decode_image_tf_example(x, input_image_size, num_channels))
+    image_tensor = tf.map_fn(
+        decode_image_tf_example_fn,
+        elems=inputs,
+        fn_output_signature=tf.TensorSpec(
+            shape=[None] * len(input_image_size) + [num_channels],
+            dtype=tf.uint8),
+    )
+  elif input_type == 'image_bytes':
+    decode_image_fn = lambda x: decode_image(x, input_image_size, num_channels)
+    image_tensor = tf.map_fn(
+        decode_image_fn, elems=inputs,
+        fn_output_signature=tf.TensorSpec(
+            shape=[None] * len(input_image_size) + [num_channels],
+            dtype=tf.uint8),)
+  else:
+    image_tensor = inputs
+  return image_tensor

orbit/utils/common.py

@@ -54,7 +54,9 @@ def make_distributed_dataset(strategy, dataset_or_fn, *args, **kwargs):
       an argument named `input_context` which will be passed a
       `tf.distribute.InputContext` instance.
     *args: Any positional arguments to pass through to `dataset_or_fn`.
-    **kwargs: Any keyword arguments to pass through to `dataset_or_fn`.
+    **kwargs: Any keyword arguments to pass through to `dataset_or_fn`, except
+      that the `input_options` keyword is used to specify a
+      `tf.distribute.InputOptions` for making the distributed dataset.
 
   Returns:
     A distributed Dataset.
@@ -62,8 +64,11 @@ def make_distributed_dataset(strategy, dataset_or_fn, *args, **kwargs):
   if strategy is None:
     strategy = tf.distribute.get_strategy()
 
+  input_options = kwargs.pop("input_options", None)
+
   if isinstance(dataset_or_fn, tf.data.Dataset):
-    return strategy.experimental_distribute_dataset(dataset_or_fn)
+    return strategy.experimental_distribute_dataset(dataset_or_fn,
+                                                    input_options)
 
   if not callable(dataset_or_fn):
     raise ValueError("`dataset_or_fn` should be either callable or an instance "
@@ -82,7 +87,7 @@ def make_distributed_dataset(strategy, dataset_or_fn, *args, **kwargs):
       kwargs["input_context"] = input_context
     return dataset_or_fn(*args, **kwargs)
 
-  return strategy.distribute_datasets_from_function(dataset_fn)
+  return strategy.distribute_datasets_from_function(dataset_fn, input_options)
 
 
 def get_value(x):

research/object_detection/data_decoders/tf_example_decoder.py

@@ -611,7 +611,6 @@ class TfExampleDecoder(data_decoder.DataDecoder):
             np.nan * tf.ones_like(tensor_dict[gt_kpt_fld]))
       else:
         num_instances = tf.shape(tensor_dict['groundtruth_classes'])[0]
-
         def true_fn(num_instances):
           """Logics to process the tensor when num_instances is not zero."""
           kpts_idx = tf.cast(self._kpts_name_to_id_table.lookup(
@@ -625,19 +624,25 @@ class TfExampleDecoder(data_decoder.DataDecoder):
               [1, num_kpt_texts])
           # Prepare the index of the keypoints to scatter the keypoint
           # coordinates: [num_kpts_texts * num_instances, 2].
-          kpt_idx = tf.concat([
+          full_kpt_idx = tf.concat([
               tf.reshape(
                   instance_idx, shape=[num_kpt_texts * num_instances, 1]),
               tf.expand_dims(kpts_idx, axis=-1)
           ], axis=1)
 
+          # Get the mask and gather only the keypoints with non-negative
+          # indices (i.e. the keypoint labels in the image/object/keypoint/text
+          # but do not exist in the label map).
+          valid_mask = tf.greater_equal(kpts_idx, 0)
+          full_kpt_idx = tf.boolean_mask(full_kpt_idx, valid_mask)
+
           gt_kpt = tf.scatter_nd(
-              kpt_idx,
-              tensor_dict[gt_kpt_fld],
+              full_kpt_idx,
+              tf.boolean_mask(tensor_dict[gt_kpt_fld], valid_mask),
               shape=[num_instances, self._num_keypoints, 2])
           gt_kpt_vis = tf.cast(tf.scatter_nd(
-              kpt_idx,
-              tensor_dict[gt_kpt_vis_fld],
+              full_kpt_idx,
+              tf.boolean_mask(tensor_dict[gt_kpt_vis_fld], valid_mask),
               shape=[num_instances, self._num_keypoints]), dtype=tf.bool)
           visibilities_tiled = tf.tile(
               tf.expand_dims(gt_kpt_vis, axis=-1), [1, 1, 2])
@@ -1091,3 +1096,4 @@ class TfExampleDecoder(data_decoder.DataDecoder):
       new_object_field = tf.repeat(
           object_field, tf.reduce_sum(expanded_indices, axis=1), axis=0)
     return new_object_field
+

research/object_detection/data_decoders/tf_example_decoder_test.py

@@ -620,6 +620,116 @@ class TfExampleDecoderTest(test_case.TestCase):
     self.assertAllEqual(
         np.zeros([2, 2], dtype=np.bool), output[gt_kpts_vis_fld][:, 3:])
 
+  def testDecodeKeypointWithKptsLabelsNotInText(self):
+    image_tensor = np.random.randint(256, size=(4, 5, 3)).astype(np.uint8)
+    encoded_jpeg, _ = self._create_encoded_and_decoded_data(
+        image_tensor, 'jpeg')
+    bbox_classes = [0, 1]
+    bbox_ymins = [0.0, 4.0]
+    bbox_xmins = [1.0, 5.0]
+    bbox_ymaxs = [2.0, 6.0]
+    bbox_xmaxs = [3.0, 7.0]
+    keypoint_ys = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
+    keypoint_xs = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
+    keypoint_visibility = [1, 2, 0, 1, 0, 2]
+    keypoint_texts = [
+        six.b('nose'), six.b('left_eye'), six.b('right_eye'), six.b('nose'),
+        six.b('left_eye'), six.b('right_eye')
+    ]
+
+    label_map_string = """
+      item: {
+        id: 1
+        name: 'face'
+        display_name: 'face'
+        keypoints {
+         id: 0
+         label: "missing_part"
+        }
+        keypoints {
+         id: 2
+         label: "right_eye"
+        }
+        keypoints {
+         id: 3
+         label: "nose"
+        }
+      }
+      item: {
+        id: 2
+        name: 'person'
+        display_name: 'person'
+        keypoints {
+         id: 1
+         label: "left_eye"
+        }
+      }
+    """
+    label_map_proto_file = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
+    with tf.gfile.Open(label_map_proto_file, 'wb') as f:
+      f.write(label_map_string)
+
+    def graph_fn():
+      example = tf.train.Example(
+          features=tf.train.Features(
+              feature={
+                  'image/encoded':
+                      dataset_util.bytes_feature(encoded_jpeg),
+                  'image/format':
+                      dataset_util.bytes_feature(six.b('jpeg')),
+                  'image/object/bbox/ymin':
+                      dataset_util.float_list_feature(bbox_ymins),
+                  'image/object/bbox/xmin':
+                      dataset_util.float_list_feature(bbox_xmins),
+                  'image/object/bbox/ymax':
+                      dataset_util.float_list_feature(bbox_ymaxs),
+                  'image/object/bbox/xmax':
+                      dataset_util.float_list_feature(bbox_xmaxs),
+                  'image/object/keypoint/y':
+                      dataset_util.float_list_feature(keypoint_ys),
+                  'image/object/keypoint/x':
+                      dataset_util.float_list_feature(keypoint_xs),
+                  'image/object/keypoint/visibility':
+                      dataset_util.int64_list_feature(keypoint_visibility),
+                  'image/object/keypoint/text':
+                      dataset_util.bytes_list_feature(keypoint_texts),
+                  'image/object/class/label':
+                      dataset_util.int64_list_feature(bbox_classes),
+              })).SerializeToString()
+
+      example_decoder = tf_example_decoder.TfExampleDecoder(
+          label_map_proto_file=label_map_proto_file, num_keypoints=5,
+          use_keypoint_label_map=True)
+      output = example_decoder.decode(tf.convert_to_tensor(example))
+
+      self.assertAllEqual((output[
+          fields.InputDataFields.groundtruth_boxes].get_shape().as_list()),
+                          [None, 4])
+      self.assertAllEqual((output[
+          fields.InputDataFields.groundtruth_keypoints].get_shape().as_list()),
+                          [None, 5, 2])
+      return output
+
+    output = self.execute_cpu(graph_fn, [])
+    expected_boxes = np.vstack([bbox_ymins, bbox_xmins, bbox_ymaxs,
+                                bbox_xmaxs]).transpose()
+    self.assertAllEqual(expected_boxes,
+                        output[fields.InputDataFields.groundtruth_boxes])
+
+    expected_keypoints = [[[np.nan, np.nan], [1., 2.], [np.nan, np.nan],
+                           [0., 1.], [np.nan, np.nan]],
+                          [[np.nan, np.nan], [np.nan, np.nan], [5., 6.],
+                           [3., 4.], [np.nan, np.nan]]]
+
+    gt_kpts_vis_fld = fields.InputDataFields.groundtruth_keypoint_visibilities
+    self.assertAllClose(expected_keypoints,
+                        output[fields.InputDataFields.groundtruth_keypoints])
+
+    expected_visibility = [[False, True, False, True, False],
+                           [False, False, True, True, False]]
+    gt_kpts_vis_fld = fields.InputDataFields.groundtruth_keypoint_visibilities
+    self.assertAllEqual(expected_visibility, output[gt_kpts_vis_fld])
+
   def testDecodeKeypointNoVisibilities(self):
     image_tensor = np.random.randint(256, size=(4, 5, 3)).astype(np.uint8)
     encoded_jpeg, _ = self._create_encoded_and_decoded_data(

research/object_detection/meta_architectures/center_net_meta_arch.py

@@ -301,8 +301,21 @@ def top_k_feature_map_locations(feature_map, max_pool_kernel_size=3, k=100,
                                                   perm=[0, 3, 1, 2])
       feature_map_peaks_transposed = tf.reshape(
           feature_map_peaks_transposed, [batch_size, num_channels, -1])
+      # safe_k will be used whenever there are fewer positions in the heatmap
+      # than the requested number of locations to score. In that case, all
+      # positions are returned in sorted order. To ensure consistent shapes for
+      # downstream ops the outputs are padded with zeros. Safe_k is also
+      # fine for TPU because TPUs require a fixed input size so the number of
+      # positions will also be fixed.
+      safe_k = tf.minimum(k, tf.shape(feature_map_peaks_transposed)[-1])
       scores, peak_flat_indices = tf.math.top_k(
-          feature_map_peaks_transposed, k=k)
+          feature_map_peaks_transposed, k=safe_k)
+      scores = tf.pad(scores, [(0, 0), (0, 0), (0, k - safe_k)])
+      peak_flat_indices = tf.pad(peak_flat_indices,
+                                 [(0, 0), (0, 0), (0, k - safe_k)])
+      scores = tf.ensure_shape(scores, (batch_size, num_channels, k))
+      peak_flat_indices = tf.ensure_shape(peak_flat_indices,
+                                          (batch_size, num_channels, k))
     # Convert the indices such that they represent the location in the full
     # (flattened) feature map of size [batch, height * width * channels].
     channel_idx = tf.range(num_channels)[tf.newaxis, :, tf.newaxis]

research/object_detection/utils/ops.py

@@ -811,7 +811,7 @@ def reframe_image_corners_relative_to_boxes(boxes):
   Returns:
     reframed_boxes: Reframes boxes with same shape as input.
   """
-  ymin, xmin, ymax, xmax = tf.unstack(boxes, axis=1)
+  ymin, xmin, ymax, xmax = (boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3])
 
   height = tf.maximum(ymax - ymin, 1e-4)
   width = tf.maximum(xmax - xmin, 1e-4)