Tensorflow/models/official/vision/dataloaders/tfexample_utils.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.

"""Utility functions to create tf.Example and tf.SequnceExample for test.

Example:video classification end-to-end test
i.e. from reading input file to train and eval.

```python
class FooTrainTest(tf.test.TestCase):

  def setUp(self):
    super(TrainTest, self).setUp()

    # Write the fake tf.train.SequenceExample to file for test.
    data_dir = os.path.join(self.get_temp_dir(), 'data')
    tf.io.gfile.makedirs(data_dir)
    self._data_path = os.path.join(data_dir, 'data.tfrecord')
    examples = [
        tfexample_utils.make_video_test_example(
            image_shape=(36, 36, 3),
            audio_shape=(20, 128),
            label=random.randint(0, 100)) for _ in range(2)
    ]
    tfexample_utils.dump_to_tfrecord(self._data_path, tf_examples=examples)

  def test_foo(self):
    dataset = tf.data.TFRecordDataset(self._data_path)
    ...

```

"""
from typing import Mapping, Optional, Sequence, Union

import numpy as np
import tensorflow as tf, tf_keras

from official.core import file_writers
from official.vision.data import fake_feature_generator
from official.vision.data import image_utils
from official.vision.data import tf_example_builder

IMAGE_KEY = 'image/encoded'
CLASSIFICATION_LABEL_KEY = 'image/class/label'
DISTILLATION_LABEL_KEY = 'image/class/soft_labels'
LABEL_KEY = 'clip/label/index'
AUDIO_KEY = 'features/audio'
DUMP_SOURCE_ID = b'7435790'


def encode_image(image_array: np.ndarray, fmt: str) -> bytes:
  return image_utils.encode_image(image_array, fmt)


def make_image_bytes(shape: Sequence[int], fmt: str = 'JPEG') -> bytes:
  """Generates image and return bytes in specified format."""
  image = fake_feature_generator.generate_image_np(*shape)
  return encode_image(image, fmt=fmt)


def put_int64_to_context(seq_example: tf.train.SequenceExample,
                         label: int = 0,
                         key: str = LABEL_KEY):
  """Puts int64 to SequenceExample context with key."""
  seq_example.context.feature[key].int64_list.value[:] = [label]


def put_bytes_list_to_feature(seq_example: tf.train.SequenceExample,
                              raw_image_bytes: bytes,
                              key: str = IMAGE_KEY,
                              repeat_num: int = 2):
  """Puts bytes list to SequenceExample context with key."""
  for _ in range(repeat_num):
    seq_example.feature_lists.feature_list.get_or_create(
        key).feature.add().bytes_list.value[:] = [raw_image_bytes]


def put_float_list_to_feature(seq_example: tf.train.SequenceExample,
                              value: Sequence[Sequence[float]], key: str):
  """Puts float list to SequenceExample context with key."""
  for s in value:
    seq_example.feature_lists.feature_list.get_or_create(
        key).feature.add().float_list.value[:] = s


def make_video_test_example(image_shape: Sequence[int] = (263, 320, 3),
                            audio_shape: Sequence[int] = (10, 256),
                            label: int = 42):
  """Generates data for testing video models (inc. RGB, audio, & label)."""
  raw_image_bytes = make_image_bytes(shape=image_shape)
  random_audio = np.random.normal(size=audio_shape).tolist()

  seq_example = tf.train.SequenceExample()
  put_int64_to_context(seq_example, label=label, key=LABEL_KEY)
  put_bytes_list_to_feature(
      seq_example, raw_image_bytes, key=IMAGE_KEY, repeat_num=4)

  put_float_list_to_feature(seq_example, value=random_audio, key=AUDIO_KEY)
  return seq_example


def dump_to_tfrecord(
    record_file: str,
    tf_examples: Sequence[Union[tf.train.Example, tf.train.SequenceExample]],
    file_type: str = 'tfrecord'):
  """Writes serialized Example to TFRecord file with path.

  Note that examples are expected to be not seriazlied.

  Args:
    record_file: The name of the output file.
    tf_examples: A list of examples to be stored.
    file_type: A string indicating the file format, could be: 'tfrecord',
      'tfrecords', 'tfrecord_compressed', 'tfrecords_gzip', 'riegeli'. The
      string is case insensitive.
  """
  file_writers.write_small_dataset(tf_examples, record_file, file_type)


def create_classification_example(
    image_height: int,
    image_width: int,
    image_format: str = 'JPEG',
    is_multilabel: bool = False,
    output_serialized_example: bool = True) -> tf.train.Example:
  """Creates image and labels for image classification input pipeline.

  Args:
    image_height: The height of test image.
    image_width: The width of test image.
    image_format: The format of test image.
    is_multilabel: A boolean flag represents whether the test image can have
      multiple labels.
    output_serialized_example: A boolean flag represents whether to return a
      serialized example.

  Returns:
    A tf.train.Example for testing.
  """
  image = fake_feature_generator.generate_image_np(image_height, image_width)
  labels = fake_feature_generator.generate_classes_np(2,
                                                      int(is_multilabel) +
                                                      1).tolist()
  builder = tf_example_builder.TfExampleBuilder()
  example = builder.add_image_matrix_feature(image, image_format,
                                             DUMP_SOURCE_ID).add_ints_feature(
                                                 CLASSIFICATION_LABEL_KEY,
                                                 labels).example
  if output_serialized_example:
    return example.SerializeToString()
  return example


def create_distillation_example(
    image_height: int,
    image_width: int,
    num_labels: int,
    image_format: str = 'JPEG',
    output_serialized_example: bool = True) -> tf.train.Example:
  """Creates image and labels for image classification with distillation.

  Args:
    image_height: The height of test image.
    image_width: The width of test image.
    num_labels: The number of labels used in test image.
    image_format: The format of test image.
    output_serialized_example: A boolean flag represents whether to return a
      serialized example.

  Returns:
    A tf.train.Example for testing.
  """
  image = fake_feature_generator.generate_image_np(image_height, image_width)
  labels = fake_feature_generator.generate_classes_np(2, 1).tolist()
  soft_labels = (fake_feature_generator.generate_classes_np(1, num_labels) +
                 0.6).tolist()
  builder = tf_example_builder.TfExampleBuilder()
  example = builder.add_image_matrix_feature(image, image_format,
                                             DUMP_SOURCE_ID).add_ints_feature(
                                                 CLASSIFICATION_LABEL_KEY,
                                                 labels).add_floats_feature(
                                                     DISTILLATION_LABEL_KEY,
                                                     soft_labels).example
  if output_serialized_example:
    return example.SerializeToString()
  return example


def create_3d_image_test_example(
    image_height: int,
    image_width: int,
    image_volume: int,
    image_channel: int,
    num_classes: int = 2,
    output_serialized_example: bool = False) -> tf.train.Example:
  """Creates 3D image and label.

  Args:
    image_height: The height of test 3D image.
    image_width: The width of test 3D image.
    image_volume: The volume of test 3D image.
    image_channel: The channel of test 3D image.
    num_classes: The number of classes of the test 3D label.
    output_serialized_example: A boolean flag represents whether to return a
      serialized example.

  Returns:
    A tf.train.Example for testing.
  """
  image = fake_feature_generator.generate_image_np(image_height, image_width,
                                                   image_channel)
  images = image[:, :, np.newaxis, :]
  images = np.tile(images, [1, 1, image_volume, 1]).astype(np.float32)

  label_shape = [image_height, image_width, image_volume, num_classes]
  labels = (
      fake_feature_generator.generate_classes_np(
          num_classes, np.prod(label_shape)
      )
      .reshape(label_shape)
      .astype(np.float32)
  )

  builder = tf_example_builder.TfExampleBuilder()
  example = builder.add_bytes_feature(IMAGE_KEY,
                                      images.tobytes()).add_bytes_feature(
                                          CLASSIFICATION_LABEL_KEY,
                                          labels.tobytes()).example
  if output_serialized_example:
    return example.SerializeToString()
  return example


def create_detection_test_example(
    image_height: int,
    image_width: int,
    image_channel: int,
    num_instances: int,
    fill_image_size: bool = True,
    output_serialized_example: bool = False) -> 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.
    fill_image_size: If image height and width will be added to the example.
    output_serialized_example: A boolean flag represents whether to return a
      serialized example.

  Returns:
    A tf.train.Example for testing.
  """
  image = fake_feature_generator.generate_image_np(image_height, image_width,
                                                   image_channel)
  boxes = fake_feature_generator.generate_normalized_boxes_np(num_instances)
  ymins, xmins, ymaxs, xmaxs = boxes.T.tolist()
  is_crowds = [0] * num_instances
  labels = fake_feature_generator.generate_classes_np(
      2, size=num_instances).tolist()
  labels_text = [b'class_1'] * num_instances
  masks = fake_feature_generator.generate_instance_masks_np(
      image_height, image_width, boxes)

  builder = tf_example_builder.TfExampleBuilder()

  example = builder.add_image_matrix_feature(
      image, image_source_id=DUMP_SOURCE_ID).add_boxes_feature(
          xmins, xmaxs, ymins, ymaxs,
          labels).add_instance_mask_matrices_feature(masks).add_ints_feature(
              'image/object/is_crowd',
              is_crowds).add_bytes_feature('image/object/class/text',
                                           labels_text).example
  if not fill_image_size:
    del example.features.feature['image/height']
    del example.features.feature['image/width']

  if output_serialized_example:
    return example.SerializeToString()
  return example


def create_segmentation_test_example(
    image_height: int,
    image_width: int,
    image_channel: int,
    output_serialized_example: bool = False,
    dense_features: Optional[Mapping[str, int]] = None) -> 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.
    output_serialized_example: A boolean flag represents whether to return a
      serialized example.
    dense_features: An optional dictionary of additional dense features, where
      the key is the prefix of the feature key in tf.Example and the value is
      the number of the channels of this feature.
  Returns:
    A tf.train.Example for testing.
  """
  image = fake_feature_generator.generate_image_np(image_height, image_width,
                                                   image_channel)
  mask = fake_feature_generator.generate_semantic_mask_np(
      image_height, image_width, 3)
  builder = tf_example_builder.TfExampleBuilder()
  builder.add_image_matrix_feature(
      image,
      image_source_id=DUMP_SOURCE_ID).add_semantic_mask_matrix_feature(mask)

  if dense_features:
    for prefix, channel in dense_features.items():
      dense_feature = fake_feature_generator.generate_semantic_mask_np(
          image_height, image_width, channel)
      builder.add_semantic_mask_matrix_feature(
          dense_feature, feature_prefix=prefix)

  example = builder.example

  if output_serialized_example:
    return example.SerializeToString()
  return example