Refactor tests for Object Detection API. (#8688)

Internal changes

--

PiperOrigin-RevId: 316837667

research/object_detection/object_detection_tutorial.ipynb → research/object_detection/colab_tutorials/object_detection_tutorial.ipynb

@@ -10,11 +10,11 @@
         "# Object Detection API Demo\n",
         "\n",
         "\u003ctable align=\"left\"\u003e\u003ctd\u003e\n",
-        "  \u003ca target=\"_blank\"  href=\"https://colab.sandbox.google.com/github/tensorflow/models/blob/master/research/object_detection/object_detection_tutorial.ipynb\"\u003e\n",
+        "  \u003ca target=\"_blank\"  href=\"https://colab.sandbox.google.com/github/tensorflow/models/blob/master/research/object_detection/colab_tutorials/colab_tutorials/object_detection_tutorial.ipynb\"\u003e\n",
         "    \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\n",
         "  \u003c/a\u003e\n",
         "\u003c/td\u003e\u003ctd\u003e\n",
-        "  \u003ca target=\"_blank\"  href=\"https://github.com/tensorflow/models/blob/master/research/object_detection/object_detection_tutorial.ipynb\"\u003e\n",
+        "  \u003ca target=\"_blank\"  href=\"https://github.com/tensorflow/models/blob/master/research/object_detection/colab_tutorials/colab_tutorials/object_detection_tutorial.ipynb\"\u003e\n",
         "    \u003cimg width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
         "\u003c/td\u003e\u003c/table\u003e"
       ]

research/object_detection/core/batcher.py

@@ -24,10 +24,6 @@ from six.moves import range
 import tensorflow.compat.v1 as tf
 
 from object_detection.core import prefetcher
-from object_detection.utils import tf_version
-
-if not tf_version.is_tf1():
-  raise ValueError('`batcher.py` is only supported in Tensorflow 1.X')
 
 rt_shape_str = '_runtime_shapes'
 

research/object_detection/core/batcher_tf1_test.py

@@ -19,14 +19,17 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import unittest
 import numpy as np
 from six.moves import range
 import tensorflow.compat.v1 as tf
 import tf_slim as slim
 
 from object_detection.core import batcher
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class BatcherTest(tf.test.TestCase):
 
   def test_batch_and_unpad_2d_tensors_of_different_sizes_in_1st_dimension(self):

research/object_detection/core/freezable_batch_norm_test.py → research/object_detection/core/freezable_batch_norm_tf2_test.py

@@ -17,15 +17,17 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
+import unittest
 import numpy as np
 from six.moves import zip
 import tensorflow.compat.v1 as tf
 
 
 from object_detection.core import freezable_batch_norm
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class FreezableBatchNormTest(tf.test.TestCase):
   """Tests for FreezableBatchNorm operations."""
 

research/object_detection/core/losses.py

@@ -681,3 +681,95 @@ class HardExampleMiner(object):
             num_positives, num_negatives)
 
 
+class PenaltyReducedLogisticFocalLoss(Loss):
+  """Penalty-reduced pixelwise logistic regression with focal loss.
+
+  The loss is defined in Equation (1) of the Objects as Points[1] paper.
+  Although the loss is defined per-pixel in the output space, this class
+  assumes that each pixel is an anchor to be compatible with the base class.
+
+  [1]: https://arxiv.org/abs/1904.07850
+  """
+
+  def __init__(self, alpha=2.0, beta=4.0, sigmoid_clip_value=1e-4):
+    """Constructor.
+
+    Args:
+      alpha: Focussing parameter of the focal loss. Increasing this will
+        decrease the loss contribution of the well classified examples.
+      beta: The local penalty reduction factor. Increasing this will decrease
+        the contribution of loss due to negative pixels near the keypoint.
+      sigmoid_clip_value: The sigmoid operation used internally will be clipped
+        between [sigmoid_clip_value, 1 - sigmoid_clip_value)
+    """
+    self._alpha = alpha
+    self._beta = beta
+    self._sigmoid_clip_value = sigmoid_clip_value
+    super(PenaltyReducedLogisticFocalLoss, self).__init__()
+
+  def _compute_loss(self, prediction_tensor, target_tensor, weights):
+    """Compute loss function.
+
+    In all input tensors, `num_anchors` is the total number of pixels in the
+    the output space.
+
+    Args:
+      prediction_tensor: A float tensor of shape [batch_size, num_anchors,
+        num_classes] representing the predicted unscaled logits for each class.
+        The function will compute sigmoid on this tensor internally.
+      target_tensor: A float tensor of shape [batch_size, num_anchors,
+        num_classes] representing a tensor with the 'splatted' keypoints,
+        possibly using a gaussian kernel. This function assumes that
+        the target is bounded between [0, 1].
+      weights: a float tensor of shape, either [batch_size, num_anchors,
+        num_classes] or [batch_size, num_anchors, 1]. If the shape is
+        [batch_size, num_anchors, 1], all the classses are equally weighted.
+
+
+    Returns:
+      loss: a float tensor of shape [batch_size, num_anchors, num_classes]
+        representing the value of the loss function.
+    """
+
+    is_present_tensor = tf.math.equal(target_tensor, 1.0)
+    prediction_tensor = tf.clip_by_value(tf.sigmoid(prediction_tensor),
+                                         self._sigmoid_clip_value,
+                                         1 - self._sigmoid_clip_value)
+
+    positive_loss = (tf.math.pow((1 - prediction_tensor), self._alpha)*
+                     tf.math.log(prediction_tensor))
+    negative_loss = (tf.math.pow((1 - target_tensor), self._beta)*
+                     tf.math.pow(prediction_tensor, self._alpha)*
+                     tf.math.log(1 - prediction_tensor))
+
+    loss = -tf.where(is_present_tensor, positive_loss, negative_loss)
+    return loss * weights
+
+
+class L1LocalizationLoss(Loss):
+  """L1 loss or absolute difference.
+
+  When used in a per-pixel manner, each pixel should be given as an anchor.
+  """
+
+  def _compute_loss(self, prediction_tensor, target_tensor, weights):
+    """Compute loss function.
+
+    Args:
+      prediction_tensor: A float tensor of shape [batch_size, num_anchors]
+        representing the (encoded) predicted locations of objects.
+      target_tensor: A float tensor of shape [batch_size, num_anchors]
+        representing the regression targets
+      weights: a float tensor of shape [batch_size, num_anchors]
+
+    Returns:
+      loss: a float tensor of shape [batch_size, num_anchors] tensor
+        representing the value of the loss function.
+    """
+    return tf.losses.absolute_difference(
+        target_tensor,
+        prediction_tensor,
+        weights=weights,
+        loss_collection=None,
+        reduction=tf.losses.Reduction.NONE
+    )

research/object_detection/core/prefetcher.py

@@ -16,10 +16,6 @@
 """Provides functions to prefetch tensors to feed into models."""
 import tensorflow.compat.v1 as tf
 
-from object_detection.utils import tf_version
-if not tf_version.is_tf1():
-  raise ValueError('`prefetcher.py` is only supported in Tensorflow 1.X')
-
 
 def prefetch(tensor_dict, capacity):
   """Creates a prefetch queue for tensors.

research/object_detection/core/prefetcher_tf1_test.py

@@ -18,16 +18,16 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import unittest
 from six.moves import range
 import tensorflow.compat.v1 as tf
-
-# pylint: disable=g-bad-import-order,
-from object_detection.core import prefetcher
 import tf_slim as slim
 
-# pylint: disable=g-bad-import-order
+from object_detection.core import prefetcher
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class PrefetcherTest(tf.test.TestCase):
   """Test class for prefetcher."""
 

research/object_detection/core/preprocessor_test.py

@@ -19,6 +19,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import unittest
 from absl.testing import parameterized
 import numpy as np
 import six
@@ -30,11 +31,12 @@ from object_detection.core import preprocessor
 from object_detection.core import preprocessor_cache
 from object_detection.core import standard_fields as fields
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 if six.PY2:
   import mock  # pylint: disable=g-import-not-at-top
 else:
-  from unittest import mock  # pylint: disable=g-import-not-at-top
+  mock = unittest.mock  # pylint: disable=g-import-not-at-top
 
 
 class PreprocessorTest(test_case.TestCase, parameterized.TestCase):
@@ -2819,6 +2821,7 @@ class PreprocessorTest(test_case.TestCase, parameterized.TestCase):
     self.assertAllEqual(images_shape, patched_images_shape)
     self.assertAllEqual(images, patched_images)
 
+  @unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
   def testAutoAugmentImage(self):
     def graph_fn():
       preprocessing_options = []

research/object_detection/dataset_tools/context_rcnn/add_context_to_examples_tf1_test.py

+# 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.
+# ==============================================================================
+"""Tests for add_context_to_examples."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import contextlib
+import datetime
+import os
+import tempfile
+import unittest
+import numpy as np
+import six
+import tensorflow.compat.v1 as tf
+
+from object_detection.dataset_tools.context_rcnn import add_context_to_examples
+from object_detection.utils import tf_version
+from apache_beam import runners
+
+
+@contextlib.contextmanager
+def InMemoryTFRecord(entries):
+  temp = tempfile.NamedTemporaryFile(delete=False)
+  filename = temp.name
+  try:
+    with tf.python_io.TFRecordWriter(filename) as writer:
+      for value in entries:
+        writer.write(value)
+    yield filename
+  finally:
+    os.unlink(temp.name)
+
+
+def BytesFeature(value):
+  return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
+
+
+def BytesListFeature(value):
+  return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
+
+
+def Int64Feature(value):
+  return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
+
+
+def Int64ListFeature(value):
+  return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
+
+
+def FloatListFeature(value):
+  return tf.train.Feature(float_list=tf.train.FloatList(value=value))
+
+
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
+class GenerateContextDataTest(tf.test.TestCase):
+
+  def _create_first_tf_example(self):
+    with self.test_session():
+      encoded_image = tf.image.encode_jpeg(
+          tf.constant(np.ones((4, 4, 3)).astype(np.uint8))).eval()
+
+    example = tf.train.Example(features=tf.train.Features(feature={
+        'image/encoded': BytesFeature(encoded_image),
+        'image/source_id': BytesFeature(six.ensure_binary('image_id_1')),
+        'image/height': Int64Feature(4),
+        'image/width': Int64Feature(4),
+        'image/object/class/label': Int64ListFeature([5, 5]),
+        'image/object/class/text': BytesListFeature([six.ensure_binary('hyena'),
+                                                     six.ensure_binary('hyena')
+                                                    ]),
+        'image/object/bbox/xmin': FloatListFeature([0.0, 0.1]),
+        'image/object/bbox/xmax': FloatListFeature([0.2, 0.3]),
+        'image/object/bbox/ymin': FloatListFeature([0.4, 0.5]),
+        'image/object/bbox/ymax': FloatListFeature([0.6, 0.7]),
+        'image/seq_id': BytesFeature(six.ensure_binary('01')),
+        'image/seq_num_frames': Int64Feature(2),
+        'image/seq_frame_num': Int64Feature(0),
+        'image/date_captured': BytesFeature(
+            six.ensure_binary(str(datetime.datetime(2020, 1, 1, 1, 0, 0)))),
+        'image/embedding': FloatListFeature([0.1, 0.2, 0.3]),
+        'image/embedding_score': FloatListFeature([0.9]),
+        'image/embedding_length': Int64Feature(3)
+
+    }))
+
+    return example.SerializeToString()
+
+  def _create_second_tf_example(self):
+    with self.test_session():
+      encoded_image = tf.image.encode_jpeg(
+          tf.constant(np.ones((4, 4, 3)).astype(np.uint8))).eval()
+
+    example = tf.train.Example(features=tf.train.Features(feature={
+        'image/encoded': BytesFeature(encoded_image),
+        'image/source_id': BytesFeature(six.ensure_binary('image_id_2')),
+        'image/height': Int64Feature(4),
+        'image/width': Int64Feature(4),
+        'image/object/class/label': Int64ListFeature([5]),
+        'image/object/class/text': BytesListFeature([six.ensure_binary('hyena')
+                                                    ]),
+        'image/object/bbox/xmin': FloatListFeature([0.0]),
+        'image/object/bbox/xmax': FloatListFeature([0.1]),
+        'image/object/bbox/ymin': FloatListFeature([0.2]),
+        'image/object/bbox/ymax': FloatListFeature([0.3]),
+        'image/seq_id': BytesFeature(six.ensure_binary('01')),
+        'image/seq_num_frames': Int64Feature(2),
+        'image/seq_frame_num': Int64Feature(1),
+        'image/date_captured': BytesFeature(
+            six.ensure_binary(str(datetime.datetime(2020, 1, 1, 1, 1, 0)))),
+        'image/embedding': FloatListFeature([0.4, 0.5, 0.6]),
+        'image/embedding_score': FloatListFeature([0.9]),
+        'image/embedding_length': Int64Feature(3)
+    }))
+
+    return example.SerializeToString()
+
+  def assert_expected_examples(self, tf_example_list):
+    self.assertAllEqual(
+        {tf_example.features.feature['image/source_id'].bytes_list.value[0]
+         for tf_example in tf_example_list},
+        {six.ensure_binary('image_id_1'), six.ensure_binary('image_id_2')})
+    self.assertAllClose(
+        tf_example_list[0].features.feature[
+            'image/context_features'].float_list.value,
+        [0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
+    self.assertAllClose(
+        tf_example_list[1].features.feature[
+            'image/context_features'].float_list.value,
+        [0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
+
+  def assert_expected_sequence_example(self, tf_sequence_example_list):
+    tf_sequence_example = tf_sequence_example_list[0]
+    num_frames = 2
+
+    self.assertAllEqual(
+        tf_sequence_example.context.feature[
+            'clip/media_id'].bytes_list.value[0], six.ensure_binary(
+                '01_0'))
+    self.assertAllClose(
+        tf_sequence_example.context.feature[
+            'image/context_features'].float_list.value,
+        [0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
+
+    seq_feature_dict = tf_sequence_example.feature_lists.feature_list
+
+    self.assertLen(
+        seq_feature_dict['image/encoded'].feature[:],
+        num_frames)
+    actual_timestamps = [
+        feature.int64_list.value[0] for feature
+        in seq_feature_dict['image/timestamp'].feature]
+    timestamps = [0, 1]
+    self.assertAllEqual(timestamps, actual_timestamps)
+
+    # First image.
+    self.assertAllClose(
+        [0.4, 0.5],
+        seq_feature_dict['region/bbox/ymin'].feature[0].float_list.value[:])
+    self.assertAllClose(
+        [0.0, 0.1],
+        seq_feature_dict['region/bbox/xmin'].feature[0].float_list.value[:])
+    self.assertAllClose(
+        [0.6, 0.7],
+        seq_feature_dict['region/bbox/ymax'].feature[0].float_list.value[:])
+    self.assertAllClose(
+        [0.2, 0.3],
+        seq_feature_dict['region/bbox/xmax'].feature[0].float_list.value[:])
+    self.assertAllEqual(
+        [six.ensure_binary('hyena'), six.ensure_binary('hyena')],
+        seq_feature_dict['region/label/string'].feature[0].bytes_list.value[:])
+
+    # Second example.
+    self.assertAllClose(
+        [0.2],
+        seq_feature_dict['region/bbox/ymin'].feature[1].float_list.value[:])
+    self.assertAllClose(
+        [0.0],
+        seq_feature_dict['region/bbox/xmin'].feature[1].float_list.value[:])
+    self.assertAllClose(
+        [0.3],
+        seq_feature_dict['region/bbox/ymax'].feature[1].float_list.value[:])
+    self.assertAllClose(
+        [0.1],
+        seq_feature_dict['region/bbox/xmax'].feature[1].float_list.value[:])
+    self.assertAllEqual(
+        [six.ensure_binary('hyena')],
+        seq_feature_dict['region/label/string'].feature[1].bytes_list.value[:])
+
+  def assert_expected_key(self, key):
+    self.assertAllEqual(key, '01')
+
+  def assert_sorted(self, example_collection):
+    example_list = list(example_collection)
+    counter = 0
+    for example in example_list:
+      frame_num = example.features.feature[
+          'image/seq_frame_num'].int64_list.value[0]
+      self.assertGreaterEqual(frame_num, counter)
+      counter = frame_num
+
+  def assert_context(self, example_collection):
+    example_list = list(example_collection)
+    for example in example_list:
+      context = example.features.feature[
+          'image/context_features'].float_list.value
+      self.assertAllClose([0.1, 0.2, 0.3, 0.4, 0.5, 0.6], context)
+
+  def assert_resized(self, example):
+    width = example.features.feature['image/width'].int64_list.value[0]
+    self.assertAllEqual(width, 2)
+    height = example.features.feature['image/height'].int64_list.value[0]
+    self.assertAllEqual(height, 2)
+
+  def assert_size(self, example):
+    width = example.features.feature['image/width'].int64_list.value[0]
+    self.assertAllEqual(width, 4)
+    height = example.features.feature['image/height'].int64_list.value[0]
+    self.assertAllEqual(height, 4)
+
+  def test_sliding_window(self):
+    example_list = ['a', 'b', 'c', 'd', 'e', 'f', 'g']
+    max_clip_length = 3
+    stride_length = 3
+    out_list = [list(i) for i in add_context_to_examples.get_sliding_window(
+        example_list, max_clip_length, stride_length)]
+    self.assertAllEqual(out_list, [['a', 'b', 'c'],
+                                   ['d', 'e', 'f'],
+                                   ['g']])
+
+  def test_rekey_data_fn(self):
+    sequence_key = 'image/seq_id'
+    time_horizon = None
+    reduce_image_size = False
+    max_dim = None
+
+    rekey_fn = add_context_to_examples.ReKeyDataFn(
+        sequence_key, time_horizon,
+        reduce_image_size, max_dim)
+    output = rekey_fn.process(self._create_first_tf_example())
+
+    self.assert_expected_key(output[0][0])
+    self.assert_size(output[0][1])
+
+  def test_rekey_data_fn_w_resize(self):
+    sequence_key = 'image/seq_id'
+    time_horizon = None
+    reduce_image_size = True
+    max_dim = 2
+
+    rekey_fn = add_context_to_examples.ReKeyDataFn(
+        sequence_key, time_horizon,
+        reduce_image_size, max_dim)
+    output = rekey_fn.process(self._create_first_tf_example())
+
+    self.assert_expected_key(output[0][0])
+    self.assert_resized(output[0][1])
+
+  def test_sort_fn(self):
+    sequence_key = 'image/seq_id'
+    sorted_image_ids = False
+    max_num_elements_in_context_features = 10
+    sort_fn = add_context_to_examples.SortGroupedDataFn(
+        sequence_key, sorted_image_ids, max_num_elements_in_context_features)
+    output = sort_fn.process(
+        ('dummy_key', [tf.train.Example.FromString(
+            self._create_second_tf_example()),
+                       tf.train.Example.FromString(
+                           self._create_first_tf_example())]))
+
+    self.assert_sorted(output[0][1])
+
+  def test_add_context_fn(self):
+    sequence_key = 'image/seq_id'
+    add_context_features = True
+    image_ids_to_keep = 'All'
+    context_fn = add_context_to_examples.GenerateContextFn(
+        sequence_key, add_context_features, image_ids_to_keep)
+    output = context_fn.process(
+        ('dummy_key', [tf.train.Example.FromString(
+            self._create_first_tf_example()),
+                       tf.train.Example.FromString(
+                           self._create_second_tf_example())]))
+
+    self.assertEqual(len(output), 2)
+    self.assert_context(output)
+
+  def test_add_context_fn_output_sequence_example(self):
+    sequence_key = 'image/seq_id'
+    add_context_features = True
+    image_ids_to_keep = 'All'
+    context_fn = add_context_to_examples.GenerateContextFn(
+        sequence_key, add_context_features, image_ids_to_keep,
+        output_type='tf_sequence_example')
+    output = context_fn.process(
+        ('01',
+         [tf.train.Example.FromString(self._create_first_tf_example()),
+          tf.train.Example.FromString(self._create_second_tf_example())]))
+
+    self.assertEqual(len(output), 1)
+    self.assert_expected_sequence_example(output)
+
+  def test_add_context_fn_output_sequence_example_cliplen(self):
+    sequence_key = 'image/seq_id'
+    add_context_features = True
+    image_ids_to_keep = 'All'
+    context_fn = add_context_to_examples.GenerateContextFn(
+        sequence_key, add_context_features, image_ids_to_keep,
+        output_type='tf_sequence_example', max_clip_length=1)
+    output = context_fn.process(
+        ('01',
+         [tf.train.Example.FromString(self._create_first_tf_example()),
+          tf.train.Example.FromString(self._create_second_tf_example())]))
+    self.assertEqual(len(output), 2)
+
+  def test_beam_pipeline(self):
+    with InMemoryTFRecord(
+        [self._create_first_tf_example(),
+         self._create_second_tf_example()]) as input_tfrecord:
+      runner = runners.DirectRunner()
+      temp_dir = tempfile.mkdtemp(dir=os.environ.get('TEST_TMPDIR'))
+      output_tfrecord = os.path.join(temp_dir, 'output_tfrecord')
+      sequence_key = six.ensure_binary('image/seq_id')
+      max_num_elements = 10
+      num_shards = 1
+      pipeline = add_context_to_examples.construct_pipeline(
+          input_tfrecord,
+          output_tfrecord,
+          sequence_key,
+          max_num_elements_in_context_features=max_num_elements,
+          num_shards=num_shards)
+      runner.run(pipeline)
+      filenames = tf.io.gfile.glob(output_tfrecord + '-?????-of-?????')
+      actual_output = []
+      record_iterator = tf.python_io.tf_record_iterator(path=filenames[0])
+      for record in record_iterator:
+        actual_output.append(record)
+      self.assertEqual(len(actual_output), 2)
+      self.assert_expected_examples([tf.train.Example.FromString(
+          tf_example) for tf_example in actual_output])
+
+  def test_beam_pipeline_sequence_example(self):
+    with InMemoryTFRecord(
+        [self._create_first_tf_example(),
+         self._create_second_tf_example()]) as input_tfrecord:
+      runner = runners.DirectRunner()
+      temp_dir = tempfile.mkdtemp(dir=os.environ.get('TEST_TMPDIR'))
+      output_tfrecord = os.path.join(temp_dir, 'output_tfrecord')
+      sequence_key = six.ensure_binary('image/seq_id')
+      max_num_elements = 10
+      num_shards = 1
+      pipeline = add_context_to_examples.construct_pipeline(
+          input_tfrecord,
+          output_tfrecord,
+          sequence_key,
+          max_num_elements_in_context_features=max_num_elements,
+          num_shards=num_shards,
+          output_type='tf_sequence_example')
+      runner.run(pipeline)
+      filenames = tf.io.gfile.glob(output_tfrecord + '-?????-of-?????')
+      actual_output = []
+      record_iterator = tf.python_io.tf_record_iterator(
+          path=filenames[0])
+      for record in record_iterator:
+        actual_output.append(record)
+      self.assertEqual(len(actual_output), 1)
+      self.assert_expected_sequence_example(
+          [tf.train.SequenceExample.FromString(
+              tf_example) for tf_example in actual_output])
+
+if __name__ == '__main__':
+  tf.test.main()

research/object_detection/dataset_tools/context_rcnn/create_cococameratraps_tfexample_tf1_test.py

+# 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.
+# ==============================================================================
+"""Tests for create_cococameratraps_tfexample_main."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import datetime
+import json
+import os
+import tempfile
+import unittest
+import numpy as np
+
+from PIL import Image
+import tensorflow.compat.v1 as tf
+from object_detection.dataset_tools.context_rcnn import create_cococameratraps_tfexample_main
+from object_detection.utils import tf_version
+from apache_beam import runners
+
+
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
+class CreateCOCOCameraTrapsTfexampleTest(tf.test.TestCase):
+
+  IMAGE_HEIGHT = 360
+  IMAGE_WIDTH = 480
+
+  def _write_random_images_to_directory(self, directory, num_frames):
+    for frame_num in range(num_frames):
+      img = np.random.randint(0, high=256,
+                              size=(self.IMAGE_HEIGHT, self.IMAGE_WIDTH, 3),
+                              dtype=np.uint8)
+      pil_image = Image.fromarray(img)
+      fname = 'im_' + str(frame_num) + '.jpg'
+      pil_image.save(os.path.join(directory, fname), 'JPEG')
+
+  def _create_json_file(self, directory, num_frames, keep_bboxes=False):
+    json_dict = {'images': [], 'annotations': []}
+    json_dict['categories'] = [{'id': 0, 'name': 'empty'},
+                               {'id': 1, 'name': 'animal'}]
+    for idx in range(num_frames):
+      im = {'id': 'im_' + str(idx),
+            'file_name': 'im_' + str(idx) + '.jpg',
+            'height': self.IMAGE_HEIGHT,
+            'width': self.IMAGE_WIDTH,
+            'seq_id': 'seq_1',
+            'seq_num_frames': num_frames,
+            'frame_num': idx,
+            'location': 'loc_' + str(idx),
+            'date_captured': str(datetime.datetime.now())
+           }
+      json_dict['images'].append(im)
+      ann = {'id': 'ann' + str(idx),
+             'image_id': 'im_' + str(idx),
+             'category_id': 1,
+            }
+      if keep_bboxes:
+        ann['bbox'] = [0.0 * self.IMAGE_WIDTH,
+                       0.1 * self.IMAGE_HEIGHT,
+                       0.5 * self.IMAGE_WIDTH,
+                       0.5 * self.IMAGE_HEIGHT]
+      json_dict['annotations'].append(ann)
+
+    json_path = os.path.join(directory, 'test_file.json')
+    with tf.io.gfile.GFile(json_path, 'w') as f:
+      json.dump(json_dict, f)
+    return json_path
+
+  def assert_expected_example_bbox(self, example):
+    self.assertAllClose(
+        example.features.feature['image/object/bbox/ymin'].float_list.value,
+        [0.1])
+    self.assertAllClose(
+        example.features.feature['image/object/bbox/xmin'].float_list.value,
+        [0.0])
+    self.assertAllClose(
+        example.features.feature['image/object/bbox/ymax'].float_list.value,
+        [0.6])
+    self.assertAllClose(
+        example.features.feature['image/object/bbox/xmax'].float_list.value,
+        [0.5])
+    self.assertAllClose(
+        example.features.feature['image/object/class/label']
+        .int64_list.value, [1])
+    self.assertAllEqual(
+        example.features.feature['image/object/class/text']
+        .bytes_list.value, ['animal'])
+    self.assertAllClose(
+        example.features.feature['image/class/label']
+        .int64_list.value, [1])
+    self.assertAllEqual(
+        example.features.feature['image/class/text']
+        .bytes_list.value, ['animal'])
+
+    # Check other essential attributes.
+    self.assertAllEqual(
+        example.features.feature['image/height'].int64_list.value,
+        [self.IMAGE_HEIGHT])
+    self.assertAllEqual(
+        example.features.feature['image/width'].int64_list.value,
+        [self.IMAGE_WIDTH])
+    self.assertAllEqual(
+        example.features.feature['image/source_id'].bytes_list.value,
+        ['im_0'])
+    self.assertTrue(
+        example.features.feature['image/encoded'].bytes_list.value)
+
+  def assert_expected_example(self, example):
+    self.assertAllClose(
+        example.features.feature['image/object/bbox/ymin'].float_list.value,
+        [])
+    self.assertAllClose(
+        example.features.feature['image/object/bbox/xmin'].float_list.value,
+        [])
+    self.assertAllClose(
+        example.features.feature['image/object/bbox/ymax'].float_list.value,
+        [])
+    self.assertAllClose(
+        example.features.feature['image/object/bbox/xmax'].float_list.value,
+        [])
+    self.assertAllClose(
+        example.features.feature['image/object/class/label']
+        .int64_list.value, [1])
+    self.assertAllEqual(
+        example.features.feature['image/object/class/text']
+        .bytes_list.value, ['animal'])
+    self.assertAllClose(
+        example.features.feature['image/class/label']
+        .int64_list.value, [1])
+    self.assertAllEqual(
+        example.features.feature['image/class/text']
+        .bytes_list.value, ['animal'])
+
+    # Check other essential attributes.
+    self.assertAllEqual(
+        example.features.feature['image/height'].int64_list.value,
+        [self.IMAGE_HEIGHT])
+    self.assertAllEqual(
+        example.features.feature['image/width'].int64_list.value,
+        [self.IMAGE_WIDTH])
+    self.assertAllEqual(
+        example.features.feature['image/source_id'].bytes_list.value,
+        ['im_0'])
+    self.assertTrue(
+        example.features.feature['image/encoded'].bytes_list.value)
+
+  def test_beam_pipeline(self):
+    runner = runners.DirectRunner()
+    num_frames = 1
+    temp_dir = tempfile.mkdtemp(dir=os.environ.get('TEST_TMPDIR'))
+    json_path = self._create_json_file(temp_dir, num_frames)
+    output_tfrecord = temp_dir+'/output'
+    self._write_random_images_to_directory(temp_dir, num_frames)
+    pipeline = create_cococameratraps_tfexample_main.create_pipeline(
+        temp_dir, json_path,
+        output_tfrecord_prefix=output_tfrecord)
+    runner.run(pipeline)
+    filenames = tf.io.gfile.glob(output_tfrecord + '-?????-of-?????')
+    actual_output = []
+    record_iterator = tf.python_io.tf_record_iterator(path=filenames[0])
+    for record in record_iterator:
+      actual_output.append(record)
+    self.assertEqual(len(actual_output), num_frames)
+    self.assert_expected_example(tf.train.Example.FromString(
+        actual_output[0]))
+
+  def test_beam_pipeline_bbox(self):
+    runner = runners.DirectRunner()
+    num_frames = 1
+    temp_dir = tempfile.mkdtemp(dir=os.environ.get('TEST_TMPDIR'))
+    json_path = self._create_json_file(temp_dir, num_frames, keep_bboxes=True)
+    output_tfrecord = temp_dir+'/output'
+    self._write_random_images_to_directory(temp_dir, num_frames)
+    pipeline = create_cococameratraps_tfexample_main.create_pipeline(
+        temp_dir, json_path,
+        output_tfrecord_prefix=output_tfrecord,
+        keep_bboxes=True)
+    runner.run(pipeline)
+    filenames = tf.io.gfile.glob(output_tfrecord+'-?????-of-?????')
+    actual_output = []
+    record_iterator = tf.python_io.tf_record_iterator(path=filenames[0])
+    for record in record_iterator:
+      actual_output.append(record)
+    self.assertEqual(len(actual_output), num_frames)
+    self.assert_expected_example_bbox(tf.train.Example.FromString(
+        actual_output[0]))
+
+
+if __name__ == '__main__':
+  tf.test.main()