Internal change

PiperOrigin-RevId: 410609558

official/vision/beta/dataloaders/tf_example_decoder_test.py

@@ -14,24 +14,13 @@
 
 """Tests for tf_example_decoder.py."""
 
-import io
 # Import libraries
 from absl.testing import parameterized
 import numpy as np
-from PIL import Image
 import tensorflow as tf
 
 from official.vision.beta.dataloaders import tf_example_decoder
-
-
-DUMP_SOURCE_ID = b'123'
-
-
-def _encode_image(image_array, fmt):
-  image = Image.fromarray(image_array)
-  with io.BytesIO() as output:
-    image.save(output, format=fmt)
-    return output.getvalue()
+from official.vision.beta.dataloaders import tfexample_utils
 
 
 class TfExampleDecoderTest(tf.test.TestCase, parameterized.TestCase):
@@ -52,73 +41,11 @@ class TfExampleDecoderTest(tf.test.TestCase, parameterized.TestCase):
     decoder = tf_example_decoder.TfExampleDecoder(
         include_mask=True, regenerate_source_id=regenerate_source_id)
 
-    image = _encode_image(
-        np.uint8(np.random.rand(image_height, image_width, 3) * 255),
-        fmt='JPEG')
-    if num_instances == 0:
-      xmins = []
-      xmaxs = []
-      ymins = []
-      ymaxs = []
-      labels = []
-      areas = []
-      is_crowds = []
-      masks = []
-    else:
-      xmins = list(np.random.rand(num_instances))
-      xmaxs = list(np.random.rand(num_instances))
-      ymins = list(np.random.rand(num_instances))
-      ymaxs = list(np.random.rand(num_instances))
-      labels = list(np.random.randint(100, size=num_instances))
-      areas = [(xmax - xmin) * (ymax - ymin) * image_height * image_width
-               for xmin, xmax, ymin, ymax in zip(xmins, xmaxs, ymins, ymaxs)]
-      is_crowds = [0] * num_instances
-      masks = []
-      for _ in range(num_instances):
-        mask = _encode_image(
-            np.uint8(np.random.rand(image_height, image_width) * 255),
-            fmt='PNG')
-        masks.append(mask)
-    serialized_example = tf.train.Example(
-        features=tf.train.Features(
-            feature={
-                'image/encoded': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=[image]))),
-                'image/source_id': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=[DUMP_SOURCE_ID]))),
-                'image/height': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=[image_height]))),
-                'image/width': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=[image_width]))),
-                'image/object/bbox/xmin': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=xmins))),
-                'image/object/bbox/xmax': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=xmaxs))),
-                'image/object/bbox/ymin': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=ymins))),
-                'image/object/bbox/ymax': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=ymaxs))),
-                'image/object/class/label': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=labels))),
-                'image/object/is_crowd': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=is_crowds))),
-                'image/object/area': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=areas))),
-                'image/object/mask': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=masks))),
-            })).SerializeToString()
+    serialized_example = tfexample_utils.create_detection_test_example(
+        image_height=image_height,
+        image_width=image_width,
+        image_channel=3,
+        num_instances=num_instances).SerializeToString()
     decoded_tensors = decoder.decode(
         tf.convert_to_tensor(value=serialized_example))
 
@@ -127,7 +54,7 @@ class TfExampleDecoderTest(tf.test.TestCase, parameterized.TestCase):
     self.assertAllEqual(
         (image_height, image_width, 3), results['image'].shape)
     if not regenerate_source_id:
-      self.assertEqual(DUMP_SOURCE_ID, results['source_id'])
+      self.assertEqual(tfexample_utils.DUMP_SOURCE_ID, results['source_id'])
     self.assertEqual(image_height, results['height'])
     self.assertEqual(image_width, results['width'])
     self.assertAllEqual(
@@ -151,7 +78,7 @@ class TfExampleDecoderTest(tf.test.TestCase, parameterized.TestCase):
                      [[0, 0, 0], [255, 255, 255], [255, 255, 255], [0, 0, 0]],
                      [[0, 0, 0], [255, 255, 255], [255, 255, 255], [0, 0, 0]],
                      [[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]]]
-    image = _encode_image(np.uint8(image_content), fmt='PNG')
+    image = tfexample_utils.encode_image(np.uint8(image_content), fmt='PNG')
     image_height = 4
     image_width = 4
     num_instances = 2
@@ -172,46 +99,38 @@ class TfExampleDecoderTest(tf.test.TestCase, parameterized.TestCase):
                      [0, 255, 255, 255],
                      [0, 255, 255, 255],
                      [0, 255, 255, 255]]]
-    masks = [_encode_image(np.uint8(m), fmt='PNG') for m in list(mask_content)]
+    masks = [
+        tfexample_utils.encode_image(np.uint8(m), fmt='PNG')
+        for m in list(mask_content)
+    ]
     serialized_example = tf.train.Example(
         features=tf.train.Features(
             feature={
-                'image/encoded': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=[image]))),
-                'image/source_id': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=[DUMP_SOURCE_ID]))),
-                'image/height': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=[image_height]))),
-                'image/width': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=[image_width]))),
-                'image/object/bbox/xmin': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=xmins))),
-                'image/object/bbox/xmax': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=xmaxs))),
-                'image/object/bbox/ymin': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=ymins))),
-                'image/object/bbox/ymax': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=ymaxs))),
-                'image/object/class/label': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=labels))),
-                'image/object/is_crowd': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=is_crowds))),
-                'image/object/area': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=areas))),
-                'image/object/mask': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=masks))),
+                'image/encoded': (tf.train.Feature(
+                    bytes_list=tf.train.BytesList(value=[image]))),
+                'image/source_id': (tf.train.Feature(
+                    bytes_list=tf.train.BytesList(
+                        value=[tfexample_utils.DUMP_SOURCE_ID]))),
+                'image/height': (tf.train.Feature(
+                    int64_list=tf.train.Int64List(value=[image_height]))),
+                'image/width': (tf.train.Feature(
+                    int64_list=tf.train.Int64List(value=[image_width]))),
+                'image/object/bbox/xmin': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=xmins))),
+                'image/object/bbox/xmax': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=xmaxs))),
+                'image/object/bbox/ymin': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=ymins))),
+                'image/object/bbox/ymax': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=ymaxs))),
+                'image/object/class/label': (tf.train.Feature(
+                    int64_list=tf.train.Int64List(value=labels))),
+                'image/object/is_crowd': (tf.train.Feature(
+                    int64_list=tf.train.Int64List(value=is_crowds))),
+                'image/object/area': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=areas))),
+                'image/object/mask': (tf.train.Feature(
+                    bytes_list=tf.train.BytesList(value=masks))),
             })).SerializeToString()
     decoded_tensors = decoder.decode(
         tf.convert_to_tensor(value=serialized_example))
@@ -221,7 +140,7 @@ class TfExampleDecoderTest(tf.test.TestCase, parameterized.TestCase):
     self.assertAllEqual(
         (image_height, image_width, 3), results['image'].shape)
     self.assertAllEqual(image_content, results['image'])
-    self.assertEqual(DUMP_SOURCE_ID, results['source_id'])
+    self.assertEqual(tfexample_utils.DUMP_SOURCE_ID, results['source_id'])
     self.assertEqual(image_height, results['height'])
     self.assertEqual(image_width, results['width'])
     self.assertAllEqual(
@@ -259,7 +178,7 @@ class TfExampleDecoderTest(tf.test.TestCase, parameterized.TestCase):
                      [[0, 0, 0], [255, 255, 255], [255, 255, 255], [0, 0, 0]],
                      [[0, 0, 0], [255, 255, 255], [255, 255, 255], [0, 0, 0]],
                      [[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]]]
-    image = _encode_image(np.uint8(image_content), fmt='PNG')
+    image = tfexample_utils.encode_image(np.uint8(image_content), fmt='PNG')
     image_height = 4
     image_width = 4
     num_instances = 2
@@ -276,40 +195,34 @@ class TfExampleDecoderTest(tf.test.TestCase, parameterized.TestCase):
                      [0, 255, 255, 255],
                      [0, 255, 255, 255],
                      [0, 255, 255, 255]]]
-    masks = [_encode_image(np.uint8(m), fmt='PNG') for m in list(mask_content)]
+    masks = [
+        tfexample_utils.encode_image(np.uint8(m), fmt='PNG')
+        for m in list(mask_content)
+    ]
     serialized_example = tf.train.Example(
         features=tf.train.Features(
             feature={
-                'image/encoded': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=[image]))),
-                'image/source_id': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=[DUMP_SOURCE_ID]))),
-                'image/height': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=[image_height]))),
-                'image/width': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=[image_width]))),
-                'image/object/bbox/xmin': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=xmins))),
-                'image/object/bbox/xmax': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=xmaxs))),
-                'image/object/bbox/ymin': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=ymins))),
-                'image/object/bbox/ymax': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=ymaxs))),
-                'image/object/class/label': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=labels))),
-                'image/object/mask': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=masks))),
+                'image/encoded': (tf.train.Feature(
+                    bytes_list=tf.train.BytesList(value=[image]))),
+                'image/source_id': (tf.train.Feature(
+                    bytes_list=tf.train.BytesList(
+                        value=[tfexample_utils.DUMP_SOURCE_ID]))),
+                'image/height': (tf.train.Feature(
+                    int64_list=tf.train.Int64List(value=[image_height]))),
+                'image/width': (tf.train.Feature(
+                    int64_list=tf.train.Int64List(value=[image_width]))),
+                'image/object/bbox/xmin': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=xmins))),
+                'image/object/bbox/xmax': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=xmaxs))),
+                'image/object/bbox/ymin': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=ymins))),
+                'image/object/bbox/ymax': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=ymaxs))),
+                'image/object/class/label': (tf.train.Feature(
+                    int64_list=tf.train.Int64List(value=labels))),
+                'image/object/mask': (tf.train.Feature(
+                    bytes_list=tf.train.BytesList(value=masks))),
             })).SerializeToString()
     decoded_tensors = decoder.decode(
         tf.convert_to_tensor(serialized_example))
@@ -318,7 +231,7 @@ class TfExampleDecoderTest(tf.test.TestCase, parameterized.TestCase):
     self.assertAllEqual(
         (image_height, image_width, 3), results['image'].shape)
     self.assertAllEqual(image_content, results['image'])
-    self.assertEqual(DUMP_SOURCE_ID, results['source_id'])
+    self.assertEqual(tfexample_utils.DUMP_SOURCE_ID, results['source_id'])
     self.assertEqual(image_height, results['height'])
     self.assertEqual(image_width, results['width'])
     self.assertAllEqual(

official/vision/beta/dataloaders/tf_example_label_map_decoder_test.py

@@ -14,28 +14,19 @@
 
 """Tests for tf_example_label_map_decoder.py."""
 
-import io
 import os
 # Import libraries
 from absl.testing import parameterized
 import numpy as np
-from PIL import Image
 import tensorflow as tf
 
 from official.vision.beta.dataloaders import tf_example_label_map_decoder
+from official.vision.beta.dataloaders import tfexample_utils
 
 
-DUMP_SOURCE_ID = b'123'
 LABEL_MAP_CSV_CONTENT = '0,class_0\n1,class_1\n2,class_2'
 
 
-def _encode_image(image_array, fmt):
-  image = Image.fromarray(image_array)
-  with io.BytesIO() as output:
-    image.save(output, format=fmt)
-    return output.getvalue()
-
-
 class TfExampleDecoderLabelMapTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.parameters(
@@ -56,74 +47,11 @@ class TfExampleDecoderLabelMapTest(tf.test.TestCase, parameterized.TestCase):
     decoder = tf_example_label_map_decoder.TfExampleDecoderLabelMap(
         label_map_path, include_mask=True)
 
-    image = _encode_image(
-        np.uint8(np.random.rand(image_height, image_width, 3) * 255),
-        fmt='JPEG')
-    if num_instances == 0:
-      xmins = []
-      xmaxs = []
-      ymins = []
-      ymaxs = []
-      labels = []
-      areas = []
-      is_crowds = []
-      masks = []
-    else:
-      xmins = list(np.random.rand(num_instances))
-      xmaxs = list(np.random.rand(num_instances))
-      ymins = list(np.random.rand(num_instances))
-      ymaxs = list(np.random.rand(num_instances))
-      labels = list(np.random.randint(100, size=num_instances))
-      areas = [(xmax - xmin) * (ymax - ymin) * image_height * image_width
-               for xmin, xmax, ymin, ymax in zip(xmins, xmaxs, ymins, ymaxs)]
-      is_crowds = [0] * num_instances
-      masks = []
-      labels = [b'class_1'] * num_instances
-      for _ in range(num_instances):
-        mask = _encode_image(
-            np.uint8(np.random.rand(image_height, image_width) * 255),
-            fmt='PNG')
-        masks.append(mask)
-    serialized_example = tf.train.Example(
-        features=tf.train.Features(
-            feature={
-                'image/encoded': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=[image]))),
-                'image/source_id': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=[DUMP_SOURCE_ID]))),
-                'image/height': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=[image_height]))),
-                'image/width': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=[image_width]))),
-                'image/object/bbox/xmin': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=xmins))),
-                'image/object/bbox/xmax': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=xmaxs))),
-                'image/object/bbox/ymin': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=ymins))),
-                'image/object/bbox/ymax': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=ymaxs))),
-                'image/object/class/text': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=labels))),
-                'image/object/is_crowd': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=is_crowds))),
-                'image/object/area': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=areas))),
-                'image/object/mask': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=masks))),
-            })).SerializeToString()
+    serialized_example = tfexample_utils.create_detection_test_example(
+        image_height=image_height,
+        image_width=image_width,
+        image_channel=3,
+        num_instances=num_instances).SerializeToString()
     decoded_tensors = decoder.decode(
         tf.convert_to_tensor(value=serialized_example))
 
@@ -131,7 +59,7 @@ class TfExampleDecoderLabelMapTest(tf.test.TestCase, parameterized.TestCase):
 
     self.assertAllEqual(
         (image_height, image_width, 3), results['image'].shape)
-    self.assertEqual(DUMP_SOURCE_ID, results['source_id'])
+    self.assertEqual(tfexample_utils.DUMP_SOURCE_ID, results['source_id'])
     self.assertEqual(image_height, results['height'])
     self.assertEqual(image_width, results['width'])
     self.assertAllEqual(
@@ -162,7 +90,7 @@ class TfExampleDecoderLabelMapTest(tf.test.TestCase, parameterized.TestCase):
                      [[0, 0, 0], [255, 255, 255], [255, 255, 255], [0, 0, 0]],
                      [[0, 0, 0], [255, 255, 255], [255, 255, 255], [0, 0, 0]],
                      [[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]]]
-    image = _encode_image(np.uint8(image_content), fmt='PNG')
+    image = tfexample_utils.encode_image(np.uint8(image_content), fmt='PNG')
     image_height = 4
     image_width = 4
     num_instances = 2
@@ -183,46 +111,38 @@ class TfExampleDecoderLabelMapTest(tf.test.TestCase, parameterized.TestCase):
                      [0, 255, 255, 255],
                      [0, 255, 255, 255],
                      [0, 255, 255, 255]]]
-    masks = [_encode_image(np.uint8(m), fmt='PNG') for m in list(mask_content)]
+    masks = [
+        tfexample_utils.encode_image(np.uint8(m), fmt='PNG')
+        for m in list(mask_content)
+    ]
     serialized_example = tf.train.Example(
         features=tf.train.Features(
             feature={
-                'image/encoded': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=[image]))),
-                'image/source_id': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=[DUMP_SOURCE_ID]))),
-                'image/height': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=[image_height]))),
-                'image/width': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=[image_width]))),
-                'image/object/bbox/xmin': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=xmins))),
-                'image/object/bbox/xmax': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=xmaxs))),
-                'image/object/bbox/ymin': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=ymins))),
-                'image/object/bbox/ymax': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=ymaxs))),
-                'image/object/class/text': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=labels))),
-                'image/object/is_crowd': (
-                    tf.train.Feature(
-                        int64_list=tf.train.Int64List(value=is_crowds))),
-                'image/object/area': (
-                    tf.train.Feature(
-                        float_list=tf.train.FloatList(value=areas))),
-                'image/object/mask': (
-                    tf.train.Feature(
-                        bytes_list=tf.train.BytesList(value=masks))),
+                'image/encoded': (tf.train.Feature(
+                    bytes_list=tf.train.BytesList(value=[image]))),
+                'image/source_id': (tf.train.Feature(
+                    bytes_list=tf.train.BytesList(
+                        value=[tfexample_utils.DUMP_SOURCE_ID]))),
+                'image/height': (tf.train.Feature(
+                    int64_list=tf.train.Int64List(value=[image_height]))),
+                'image/width': (tf.train.Feature(
+                    int64_list=tf.train.Int64List(value=[image_width]))),
+                'image/object/bbox/xmin': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=xmins))),
+                'image/object/bbox/xmax': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=xmaxs))),
+                'image/object/bbox/ymin': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=ymins))),
+                'image/object/bbox/ymax': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=ymaxs))),
+                'image/object/class/text': (tf.train.Feature(
+                    bytes_list=tf.train.BytesList(value=labels))),
+                'image/object/is_crowd': (tf.train.Feature(
+                    int64_list=tf.train.Int64List(value=is_crowds))),
+                'image/object/area': (tf.train.Feature(
+                    float_list=tf.train.FloatList(value=areas))),
+                'image/object/mask': (tf.train.Feature(
+                    bytes_list=tf.train.BytesList(value=masks))),
             })).SerializeToString()
     decoded_tensors = decoder.decode(
         tf.convert_to_tensor(value=serialized_example))
@@ -232,7 +152,7 @@ class TfExampleDecoderLabelMapTest(tf.test.TestCase, parameterized.TestCase):
     self.assertAllEqual(
         (image_height, image_width, 3), results['image'].shape)
     self.assertAllEqual(image_content, results['image'])
-    self.assertEqual(DUMP_SOURCE_ID, results['source_id'])
+    self.assertEqual(tfexample_utils.DUMP_SOURCE_ID, results['source_id'])
     self.assertEqual(image_height, results['height'])
     self.assertEqual(image_width, results['width'])
     self.assertAllEqual(

official/vision/beta/dataloaders/tfds_factory_test.py

@@ -23,6 +23,22 @@ from official.vision.beta.dataloaders import tfds_factory
 
 class TFDSFactoryTest(tf.test.TestCase, parameterized.TestCase):
 
+  def _create_test_example(self):
+    serialized_example = {
+        'image': tf.ones(shape=(100, 100, 3), dtype=tf.uint8),
+        'label': 1,
+        'image/id': 0,
+        'objects': {
+            'label': 1,
+            'is_crowd': 0,
+            'area': 0.5,
+            'bbox': [0.1, 0.2, 0.3, 0.4]
+        },
+        'segmentation_label': tf.ones((100, 100, 1), dtype=tf.uint8),
+        'image_left': tf.ones(shape=(100, 100, 3), dtype=tf.uint8)
+    }
+    return serialized_example
+
   @parameterized.parameters(
       ('imagenet2012'),
       ('cifar10'),
@@ -31,6 +47,10 @@ class TFDSFactoryTest(tf.test.TestCase, parameterized.TestCase):
   def test_classification_decoder(self, tfds_name):
     decoder = tfds_factory.get_classification_decoder(tfds_name)
     self.assertIsInstance(decoder, base_decoder.Decoder)
+    decoded_tensor = decoder.decode(self._create_test_example())
+    self.assertLen(decoded_tensor, 2)
+    self.assertIn('image/encoded', decoded_tensor)
+    self.assertIn('image/class/label', decoded_tensor)
 
   @parameterized.parameters(
       ('flowers'),
@@ -48,6 +68,16 @@ class TFDSFactoryTest(tf.test.TestCase, parameterized.TestCase):
   def test_detection_decoder(self, tfds_name):
     decoder = tfds_factory.get_detection_decoder(tfds_name)
     self.assertIsInstance(decoder, base_decoder.Decoder)
+    decoded_tensor = decoder.decode(self._create_test_example())
+    self.assertLen(decoded_tensor, 8)
+    self.assertIn('image', decoded_tensor)
+    self.assertIn('source_id', decoded_tensor)
+    self.assertIn('height', decoded_tensor)
+    self.assertIn('width', decoded_tensor)
+    self.assertIn('groundtruth_classes', decoded_tensor)
+    self.assertIn('groundtruth_is_crowd', decoded_tensor)
+    self.assertIn('groundtruth_area', decoded_tensor)
+    self.assertIn('groundtruth_boxes', decoded_tensor)
 
   @parameterized.parameters(
       ('pascal'),
@@ -65,6 +95,12 @@ class TFDSFactoryTest(tf.test.TestCase, parameterized.TestCase):
   def test_segmentation_decoder(self, tfds_name):
     decoder = tfds_factory.get_segmentation_decoder(tfds_name)
     self.assertIsInstance(decoder, base_decoder.Decoder)
+    decoded_tensor = decoder.decode(self._create_test_example())
+    self.assertLen(decoded_tensor, 4)
+    self.assertIn('image/encoded', decoded_tensor)
+    self.assertIn('image/segmentation/class/encoded', decoded_tensor)
+    self.assertIn('image/height', decoded_tensor)
+    self.assertIn('image/width', decoded_tensor)
 
   @parameterized.parameters(
       ('coco'),

official/vision/beta/dataloaders/tfexample_utils.py

@@ -54,16 +54,20 @@ IMAGE_KEY = 'image/encoded'
 CLASSIFICATION_LABEL_KEY = 'image/class/label'
 LABEL_KEY = 'clip/label/index'
 AUDIO_KEY = 'features/audio'
+DUMP_SOURCE_ID = b'123'
 
 
-def make_image_bytes(shape: Sequence[int]):
-  """Generates image and return bytes in JPEG format."""
+def encode_image(image_array: np.array, fmt: str) -> bytes:
+  image = Image.fromarray(image_array)
+  with io.BytesIO() as output:
+    image.save(output, format=fmt)
+    return output.getvalue()
+
+
+def make_image_bytes(shape: Sequence[int], fmt: str = 'JPEG') -> bytes:
+  """Generates image and return bytes in specified format."""
   random_image = np.random.randint(0, 256, size=shape, dtype=np.uint8)
-  random_image = Image.fromarray(random_image)
-  with io.BytesIO() as buffer:
-    random_image.save(buffer, format='JPEG')
-    raw_image_bytes = buffer.getvalue()
-  return raw_image_bytes
+  return encode_image(random_image, fmt=fmt)
 
 
 def put_int64_to_context(seq_example: tf.train.SequenceExample,
@@ -164,3 +168,102 @@ def create_3d_image_test_example(image_height: int, image_width: int,
           bytes_list=tf.train.BytesList(value=[labels.tobytes()])))
   }
   return tf.train.Example(features=tf.train.Features(feature=feature))
+
+
+def create_detection_test_example(image_height: int, image_width: int,
+                                  image_channel: int,
+                                  num_instances: int) -> tf.train.Example:
+  """Creates and returns a test example containing box and mask annotations.
+
+  Args:
+    image_height: The height of test image.
+    image_width: The width of test image.
+    image_channel: The channel of test image.
+    num_instances: The number of object instances per image.
+
+  Returns:
+    A tf.train.Example for testing.
+  """
+  image = make_image_bytes([image_height, image_width, image_channel])
+  if num_instances == 0:
+    xmins = []
+    xmaxs = []
+    ymins = []
+    ymaxs = []
+    labels = []
+    areas = []
+    is_crowds = []
+    masks = []
+    labels_text = []
+  else:
+    xmins = list(np.random.rand(num_instances))
+    xmaxs = list(np.random.rand(num_instances))
+    ymins = list(np.random.rand(num_instances))
+    ymaxs = list(np.random.rand(num_instances))
+    labels_text = [b'class_1'] * num_instances
+    labels = list(np.random.randint(100, size=num_instances))
+    areas = [(xmax - xmin) * (ymax - ymin) * image_height * image_width
+             for xmin, xmax, ymin, ymax in zip(xmins, xmaxs, ymins, ymaxs)]
+    is_crowds = [0] * num_instances
+    masks = []
+    for _ in range(num_instances):
+      mask = make_image_bytes([image_height, image_width], fmt='PNG')
+      masks.append(mask)
+  return tf.train.Example(
+      features=tf.train.Features(
+          feature={
+              'image/encoded': (tf.train.Feature(
+                  bytes_list=tf.train.BytesList(value=[image]))),
+              'image/source_id': (tf.train.Feature(
+                  bytes_list=tf.train.BytesList(value=[DUMP_SOURCE_ID]))),
+              'image/height': (tf.train.Feature(
+                  int64_list=tf.train.Int64List(value=[image_height]))),
+              'image/width': (tf.train.Feature(
+                  int64_list=tf.train.Int64List(value=[image_width]))),
+              'image/object/bbox/xmin': (tf.train.Feature(
+                  float_list=tf.train.FloatList(value=xmins))),
+              'image/object/bbox/xmax': (tf.train.Feature(
+                  float_list=tf.train.FloatList(value=xmaxs))),
+              'image/object/bbox/ymin': (tf.train.Feature(
+                  float_list=tf.train.FloatList(value=ymins))),
+              'image/object/bbox/ymax': (tf.train.Feature(
+                  float_list=tf.train.FloatList(value=ymaxs))),
+              'image/object/class/label': (tf.train.Feature(
+                  int64_list=tf.train.Int64List(value=labels))),
+              'image/object/class/text': (tf.train.Feature(
+                  bytes_list=tf.train.BytesList(value=labels_text))),
+              'image/object/is_crowd': (tf.train.Feature(
+                  int64_list=tf.train.Int64List(value=is_crowds))),
+              'image/object/area': (tf.train.Feature(
+                  float_list=tf.train.FloatList(value=areas))),
+              'image/object/mask': (tf.train.Feature(
+                  bytes_list=tf.train.BytesList(value=masks))),
+          }))
+
+
+def create_segmentation_test_example(image_height: int, image_width: int,
+                                     image_channel: int) -> tf.train.Example:
+  """Creates and returns a test example containing mask annotations.
+
+  Args:
+    image_height: The height of test image.
+    image_width: The width of test image.
+    image_channel: The channel of test image.
+
+  Returns:
+    A tf.train.Example for testing.
+  """
+  image = make_image_bytes([image_height, image_width, image_channel])
+  mask = make_image_bytes([image_height, image_width], fmt='PNG')
+  return tf.train.Example(
+      features=tf.train.Features(
+          feature={
+              'image/encoded': (tf.train.Feature(
+                  bytes_list=tf.train.BytesList(value=[image]))),
+              'image/segmentation/class/encoded': (tf.train.Feature(
+                  bytes_list=tf.train.BytesList(value=[mask]))),
+              'image/height': (tf.train.Feature(
+                  int64_list=tf.train.Int64List(value=[image_height]))),
+              'image/width': (tf.train.Feature(
+                  int64_list=tf.train.Int64List(value=[image_width])))
+          }))