models/official/vision/data/tf_example_builder_test.py

# Copyright 2023 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 tf_example_builder."""

from absl.testing import parameterized
import tensorflow as tf, tf_keras
from official.vision.data import fake_feature_generator
from official.vision.data import image_utils
from official.vision.data import tf_example_builder


class TfExampleBuilderTest(tf.test.TestCase, parameterized.TestCase):

  @parameterized.named_parameters(('RGB_PNG', 128, 64, 3, 'PNG', 3),
                                  ('RGB_RAW', 128, 128, 3, 'RAW', 0),
                                  ('RGB_JPEG', 64, 128, 3, 'JPEG', [2, 5]))
  def test_add_image_matrix_feature_success(self, height, width, num_channels,
                                            image_format, label):
    # Prepare test data.
    image_np = fake_feature_generator.generate_image_np(height, width,
                                                        num_channels)
    expected_image_bytes = image_utils.encode_image(image_np, image_format)
    hashed_image = bytes('10242048', 'ascii')

    # Run code logic.
    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_image_matrix_feature(
        image_np, image_format, hashed_image, label=label)
    example = example_builder.example

    # Verify outputs.
    # Prefer to use string literal for feature keys to directly display the
    # structure of the expected tf.train.Example.
    if isinstance(label, int):
      expected_labels = [label]
    else:
      expected_labels = label
    self.assertProtoEquals(
        tf.train.Example(
            features=tf.train.Features(
                feature={
                    'image/encoded':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[expected_image_bytes])),
                    'image/format':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[bytes(image_format, 'ascii')])),
                    'image/height':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(value=[height])),
                    'image/width':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(value=[width])),
                    'image/channels':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(
                                value=[num_channels])),
                    'image/source_id':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[hashed_image])),
                    'image/class/label':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(
                                value=expected_labels)),
                })), example)

  def test_add_image_matrix_feature_with_feature_prefix_success(self):
    height = 64
    width = 64
    num_channels = 1
    image_format = 'PNG'
    feature_prefix = 'depth'
    label = 8
    image_np = fake_feature_generator.generate_image_np(height, width,
                                                        num_channels)
    expected_image_bytes = image_utils.encode_image(image_np, image_format)
    hashed_image = bytes('10242048', 'ascii')

    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_image_matrix_feature(
        image_np,
        image_format,
        hashed_image,
        feature_prefix=feature_prefix,
        label=label)
    example = example_builder.example

    self.assertProtoEquals(
        tf.train.Example(
            features=tf.train.Features(
                feature={
                    'depth/image/encoded':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[expected_image_bytes])),
                    'depth/image/format':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[bytes(image_format, 'ascii')])),
                    'depth/image/height':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(value=[height])),
                    'depth/image/width':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(value=[width])),
                    'depth/image/channels':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(
                                value=[num_channels])),
                    'depth/image/source_id':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[hashed_image])),
                    'depth/image/class/label':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(value=[label]))
                })), example)

  def test_add_encoded_raw_image_feature_success(self):
    height = 128
    width = 128
    num_channels = 3
    image_format = 'RAW'
    expected_image_bytes = bytes('image', 'ascii')
    hashed_image = bytes('16188651', 'ascii')

    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_encoded_image_feature(expected_image_bytes, 'RAW',
                                              height, width, num_channels)
    example = example_builder.example

    self.assertProtoEquals(
        tf.train.Example(
            features=tf.train.Features(
                feature={
                    'image/encoded':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[expected_image_bytes])),
                    'image/format':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[bytes(image_format, 'ascii')])),
                    'image/height':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(value=[height])),
                    'image/width':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(value=[width])),
                    'image/channels':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(
                                value=[num_channels])),
                    'image/source_id':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(value=[hashed_image]))
                })), example)

  def test_add_encoded_raw_image_feature_valueerror(self):
    image_format = 'RAW'
    image_bytes = tf.bfloat16.as_numpy_dtype
    image_np = fake_feature_generator.generate_image_np(1, 1, 1)
    image_np = image_np.astype(image_bytes)
    expected_image_bytes = image_utils.encode_image(image_np, image_format)

    example_builder = tf_example_builder.TfExampleBuilder()
    with self.assertRaises(ValueError):
      example_builder.add_encoded_image_feature(expected_image_bytes,
                                                image_format)

  @parameterized.product(
      miss_image_format=(True, False),
      miss_height=(True, False),
      miss_width=(True, False),
      miss_num_channels=(True, False),
      miss_label=(True, False))
  def test_add_encoded_image_feature_success(self, miss_image_format,
                                             miss_height, miss_width,
                                             miss_num_channels,
                                             miss_label):
    height = 64
    width = 64
    num_channels = 3
    image_format = 'PNG'
    image_np = fake_feature_generator.generate_image_np(height, width,
                                                        num_channels)
    image_bytes = image_utils.encode_image(image_np, image_format)
    # We don't test on image_source_id because encoding process becomes
    # non-deterministic.
    hashed_image = bytes('10242048', 'ascii')
    label = 5

    image_format = None if miss_image_format else image_format
    height = None if miss_height else height
    width = None if miss_width else width
    num_channels = None if miss_num_channels else num_channels
    label = None if miss_label else label

    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_encoded_image_feature(
        image_bytes,
        image_format=image_format,
        height=height,
        width=width,
        num_channels=num_channels,
        image_source_id=hashed_image,
        label=label)
    example = example_builder.example

    expected_features = {
        'image/encoded':
            tf.train.Feature(
                bytes_list=tf.train.BytesList(value=[image_bytes])),
        'image/format':
            tf.train.Feature(
                bytes_list=tf.train.BytesList(
                    value=[bytes('PNG', 'ascii')])),
        'image/height':
            tf.train.Feature(
                int64_list=tf.train.Int64List(value=[64])),
        'image/width':
            tf.train.Feature(
                int64_list=tf.train.Int64List(value=[64])),
        'image/channels':
            tf.train.Feature(
                int64_list=tf.train.Int64List(value=[3])),
        'image/source_id':
            tf.train.Feature(
                bytes_list=tf.train.BytesList(value=[hashed_image]))}
    if not miss_label:
      expected_features.update({
          'image/class/label':
              tf.train.Feature(
                  int64_list=tf.train.Int64List(value=[label]))})
    self.assertProtoEquals(
        tf.train.Example(features=tf.train.Features(feature=expected_features)),
        example)

  @parameterized.named_parameters(('no_box', 0), ('10_boxes', 10))
  def test_add_normalized_boxes_feature(self, num_boxes):
    normalized_boxes_np = fake_feature_generator.generate_normalized_boxes_np(
        num_boxes)
    ymins, xmins, ymaxs, xmaxs = normalized_boxes_np.T.tolist()
    labels = fake_feature_generator.generate_classes_np(
        2, size=num_boxes).tolist()

    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_boxes_feature(
        xmins, xmaxs, ymins, ymaxs, labels=labels, normalized=True)
    example = example_builder.example

    self.assertProtoEquals(
        tf.train.Example(
            features=tf.train.Features(
                feature={
                    'image/object/bbox/xmin':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(value=xmins)),
                    'image/object/bbox/ymin':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(value=ymins)),
                    'image/object/bbox/xmax':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(value=xmaxs)),
                    '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)),
                })), example)

  @parameterized.named_parameters(('no_box', 0), ('10_boxes', 10))
  def test_add_box_pixels_feature(self, num_boxes):
    height, width = 10, 10
    boxes_np = fake_feature_generator.generate_boxes_np(height, width,
                                                        num_boxes)
    ymins, xmins, ymaxs, xmaxs = boxes_np.T.tolist()
    labels = fake_feature_generator.generate_classes_np(
        2, size=num_boxes).tolist()

    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_boxes_feature(
        xmins, xmaxs, ymins, ymaxs, labels=labels, normalized=False)
    example = example_builder.example

    self.assertProtoEquals(
        tf.train.Example(
            features=tf.train.Features(
                feature={
                    'image/object/bbox/xmin_pixels':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(value=xmins)),
                    'image/object/bbox/ymin_pixels':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(value=ymins)),
                    'image/object/bbox/xmax_pixels':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(value=xmaxs)),
                    'image/object/bbox/ymax_pixels':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(value=ymaxs)),
                    'image/object/class/label':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(value=labels)),
                })), example)

  @parameterized.named_parameters(('no_box', 0), ('10_boxes', 10))
  def test_add_normalized_boxes_feature_with_confidence_and_prefix(
      self, num_boxes):
    normalized_boxes_np = fake_feature_generator.generate_normalized_boxes_np(
        num_boxes)
    ymins, xmins, ymaxs, xmaxs = normalized_boxes_np.T.tolist()
    labels = fake_feature_generator.generate_classes_np(
        2, size=num_boxes).tolist()
    confidences = fake_feature_generator.generate_confidences_np(
        size=num_boxes).tolist()
    feature_prefix = 'predicted'

    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_boxes_feature(
        xmins,
        xmaxs,
        ymins,
        ymaxs,
        labels=labels,
        confidences=confidences,
        normalized=True,
        feature_prefix=feature_prefix)
    example = example_builder.example

    self.assertProtoEquals(
        tf.train.Example(
            features=tf.train.Features(
                feature={
                    'predicted/image/object/bbox/xmin':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(value=xmins)),
                    'predicted/image/object/bbox/ymin':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(value=ymins)),
                    'predicted/image/object/bbox/xmax':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(value=xmaxs)),
                    'predicted/image/object/bbox/ymax':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(value=ymaxs)),
                    'predicted/image/object/class/label':
                        tf.train.Feature(
                            int64_list=tf.train.Int64List(value=labels)),
                    'predicted/image/object/bbox/confidence':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(value=confidences)),
                })), example)

  @parameterized.named_parameters(('no_mask', 128, 64, 0),
                                  ('10_masks', 64, 128, 10))
  def test_add_instance_mask_matrices_feature_success(self, height, width,
                                                      num_masks):
    # Prepare test data.
    instance_masks_np = fake_feature_generator.generate_instance_masks_np(
        height,
        width,
        boxes_np=fake_feature_generator.generate_boxes_np(
            height, width, num_masks),
        normalized=False)
    expected_instance_masks_bytes = list(
        map(lambda x: image_utils.encode_image(x, 'PNG'), instance_masks_np))

    # Run code logic.
    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_instance_mask_matrices_feature(instance_masks_np)
    example = example_builder.example

    # Verify outputs.
    # Prefer to use string literal for feature keys to directly display the
    # structure of the expected tf.train.Example.
    self.assertProtoEquals(
        tf.train.Example(
            features=tf.train.Features(
                feature={
                    'image/object/mask':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=expected_instance_masks_bytes)),
                    'image/object/area':
                        # The box area is 4x smaller than the image, and the
                        # mask area is 2x smaller than the box.
                        tf.train.Feature(
                            float_list=tf.train.FloatList(
                                value=[height * width / 8] * num_masks))
                })), example)

  @parameterized.named_parameters(('with_mask_areas', True),
                                  ('without_mask_areas', False))
  def test_add_encoded_instance_masks_feature_success(self, has_mask_areas):
    height = 64
    width = 64
    image_format = 'PNG'
    mask_np = fake_feature_generator.generate_semantic_mask_np(height, width, 2)
    mask_bytes = image_utils.encode_image(mask_np, image_format)

    test_masks = [mask_bytes for _ in range(2)]
    mask_areas = [2040., 2040.] if has_mask_areas else None

    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_encoded_instance_masks_feature(
        test_masks, mask_areas=mask_areas)
    example = example_builder.example

    self.assertProtoEquals(
        tf.train.Example(
            features=tf.train.Features(
                feature={
                    'image/object/mask':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(value=test_masks)),
                    'image/object/area':
                        tf.train.Feature(
                            float_list=tf.train.FloatList(
                                value=[2040., 2040.])),
                })), example)

  @parameterized.named_parameters(
      ('with_visualization_mask', 128, 64, True),
      ('without_visualization_mask', 64, 128, False))
  def test_add_semantic_mask_matrices_feature_success(self, height, width,
                                                      has_visualization_mask):
    # Prepare test data.
    semantic_mask_np = fake_feature_generator.generate_semantic_mask_np(
        height, width, 2)
    image_format = 'PNG'
    expected_feature_dict = {
        'image/segmentation/class/encoded':
            tf.train.Feature(
                bytes_list=tf.train.BytesList(value=[
                    image_utils.encode_image(semantic_mask_np, image_format)
                ])),
        'image/segmentation/class/format':
            tf.train.Feature(
                bytes_list=tf.train.BytesList(
                    value=[bytes(image_format, 'ascii')])),
    }
    visualization_mask_np = None
    if has_visualization_mask:
      visualization_mask_np = fake_feature_generator.generate_image_np(
          height, width)
      expected_feature_dict.update({
          'image/segmentation/class/visualization/encoded':
              tf.train.Feature(
                  bytes_list=tf.train.BytesList(value=[
                      image_utils.encode_image(visualization_mask_np,
                                               image_format)
                  ])),
          'image/segmentation/class/visualization/format':
              tf.train.Feature(
                  bytes_list=tf.train.BytesList(
                      value=[bytes(image_format, 'ascii')])),
      })

    # Run code logic.
    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_semantic_mask_matrix_feature(semantic_mask_np,
                                                     image_format,
                                                     visualization_mask_np,
                                                     image_format)
    example = example_builder.example

    self.assertProtoEquals(
        tf.train.Example(
            features=tf.train.Features(feature=expected_feature_dict)), example)

  @parameterized.named_parameters(('with_visualization_mask', True),
                                  ('without_visualization_mask', False))
  def test_add_encoded_semantic_mask_feature_success(self,
                                                     has_visualization_mask):
    height, width = 64, 64
    semantic_mask_np = fake_feature_generator.generate_semantic_mask_np(
        height, width, 2)
    image_format = 'PNG'
    encoded_semantic_mask = image_utils.encode_image(semantic_mask_np,
                                                     image_format)
    expected_feature_dict = {
        'image/segmentation/class/encoded':
            tf.train.Feature(
                bytes_list=tf.train.BytesList(value=[encoded_semantic_mask])),
        'image/segmentation/class/format':
            tf.train.Feature(
                bytes_list=tf.train.BytesList(
                    value=[bytes(image_format, 'ascii')])),
    }
    encoded_visualization_mask = None
    if has_visualization_mask:
      visualization_mask_np = fake_feature_generator.generate_image_np(
          height, width)
      encoded_visualization_mask = image_utils.encode_image(
          visualization_mask_np, image_format)
      expected_feature_dict.update({
          'image/segmentation/class/visualization/encoded':
              tf.train.Feature(
                  bytes_list=tf.train.BytesList(
                      value=[encoded_visualization_mask])),
          'image/segmentation/class/visualization/format':
              tf.train.Feature(
                  bytes_list=tf.train.BytesList(
                      value=[bytes(image_format, 'ascii')])),
      })

    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_encoded_semantic_mask_feature(
        encoded_semantic_mask, image_format, encoded_visualization_mask,
        image_format)
    example = example_builder.example

    self.assertProtoEquals(
        tf.train.Example(
            features=tf.train.Features(feature=expected_feature_dict)), example)

  def test_add_panoptic_mask_matrices_feature_success(self):
    # Prepare test data.
    height, width, num_instances = 64, 64, 10
    num_thing_classes, num_semantic_segmentation_classes = 3, 6
    image_format = 'PNG'

    normalized_boxes_np = fake_feature_generator.generate_normalized_boxes_np(
        num_instances)
    instance_masks_np = fake_feature_generator.generate_instance_masks_np(
        height, width, normalized_boxes_np)
    instance_classes_np = fake_feature_generator.generate_classes_np(
        num_thing_classes, num_instances)
    semantic_mask_np = fake_feature_generator.generate_semantic_mask_np(
        height, width, num_semantic_segmentation_classes)
    panoptic_category_mask_np, panoptic_instance_mask_np = (
        fake_feature_generator.generate_panoptic_masks_np(
            semantic_mask_np, instance_masks_np, instance_classes_np,
            num_thing_classes - 1))

    # Run code logic.
    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_panoptic_mask_matrix_feature(panoptic_category_mask_np,
                                                     panoptic_instance_mask_np,
                                                     image_format,
                                                     image_format)
    example = example_builder.example

    self.assertProtoEquals(
        tf.train.Example(
            features=tf.train.Features(
                feature={
                    'image/panoptic/category/encoded':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(value=[
                                image_utils.encode_image(
                                    panoptic_category_mask_np, image_format)
                            ])),
                    'image/panoptic/category/format':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[bytes(image_format, 'ascii')])),
                    'image/panoptic/instance/encoded':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(value=[
                                image_utils.encode_image(
                                    panoptic_instance_mask_np, image_format)
                            ])),
                    'image/panoptic/instance/format':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[bytes(image_format, 'ascii')])),
                })), example)

  def test_add_encoded_panoptic_mask_feature_success(self):
    # Prepare test data.
    height, width, num_instances = 64, 64, 10
    num_thing_classes, num_semantic_segmentation_classes = 3, 6
    image_format = 'PNG'

    normalized_boxes_np = fake_feature_generator.generate_normalized_boxes_np(
        num_instances)
    instance_masks_np = fake_feature_generator.generate_instance_masks_np(
        height, width, normalized_boxes_np)
    instance_classes_np = fake_feature_generator.generate_classes_np(
        num_thing_classes, num_instances)
    semantic_mask_np = fake_feature_generator.generate_semantic_mask_np(
        height, width, num_semantic_segmentation_classes)
    panoptic_category_mask_np, panoptic_instance_mask_np = (
        fake_feature_generator.generate_panoptic_masks_np(
            semantic_mask_np, instance_masks_np, instance_classes_np,
            num_thing_classes - 1))

    encoded_panoptic_category_mask = image_utils.encode_image(
        panoptic_category_mask_np, image_format)
    encoded_panoptic_instance_mask = image_utils.encode_image(
        panoptic_instance_mask_np, image_format)

    example_builder = tf_example_builder.TfExampleBuilder()
    example_builder.add_encoded_panoptic_mask_feature(
        encoded_panoptic_category_mask, encoded_panoptic_instance_mask,
        image_format, image_format)
    example = example_builder.example

    self.assertProtoEquals(
        tf.train.Example(
            features=tf.train.Features(
                feature={
                    'image/panoptic/category/encoded':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[encoded_panoptic_category_mask])),
                    'image/panoptic/category/format':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[bytes(image_format, 'ascii')])),
                    'image/panoptic/instance/encoded':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[encoded_panoptic_instance_mask])),
                    'image/panoptic/instance/format':
                        tf.train.Feature(
                            bytes_list=tf.train.BytesList(
                                value=[bytes(image_format, 'ascii')])),
                })), example)


if __name__ == '__main__':
  tf.test.main()