official/vision/data/tf_example_builder.py
# Copyright 2024 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.
"""Builder class for preparing tf.train.Example in vision tasks."""
# https://www.python.org/dev/peps/pep-0563/#enabling-the-future-behavior-in-python-3-7
from __future__ import annotations
import hashlib
from typing import Optional, Sequence, Union
import numpy as np
from official.core import tf_example_builder
from official.vision.data import image_utils
from official.vision.data import tf_example_feature_key
BytesValueType = Union[bytes, Sequence[bytes], str, Sequence[str]]
_to_array = lambda v: [v] if not isinstance(v, (list, np.ndarray)) else v
_to_bytes = lambda v: v.encode() if isinstance(v, str) else v
_to_bytes_array = lambda v: list(map(_to_bytes, _to_array(v)))
class TfExampleBuilder(tf_example_builder.TfExampleBuilder):
"""Builder class for preparing tf.train.Example in vision task.
Read API doc at https://www.tensorflow.org/api_docs/python/tf/train/Example.
"""
def add_image_matrix_feature(
self,
image_matrix: np.ndarray,
image_format: str = 'PNG',
image_source_id: Optional[bytes] = None,
feature_prefix: Optional[str] = None,
label: Optional[Union[int, Sequence[int]]] = None) -> 'TfExampleBuilder':
"""Encodes and adds image features to the example.
See `tf_example_feature_key.EncodedImageFeatureKey` for list of feature keys
that will be added to the example.
Example usages:
>>> example_builder = TfExampleBuilder()
* For adding RGB image feature with PNG encoding:
>>> example_builder.add_image_matrix_feature(image_matrix)
* For adding RGB image feature with a pre-generated source ID.
>>> example_builder.add_image_matrix_feature(
image_matrix, image_source_id=image_source_id)
* For adding single-channel depth image feature with JPEG encoding:
>>> example_builder.add_image_matrix_feature(
image_matrix, image_format=ImageFormat.JPEG,
feature_prefix='depth')
Args:
image_matrix: Numpy image matrix with shape (height, width, channels)
image_format: Image format string, defaults to 'PNG'.
image_source_id: Unique string ID to identify the image. Hashed image will
be used if the field is not provided.
feature_prefix: Feature prefix for image features.
label: the label or a list of labels for the image.
Returns:
The builder object for subsequent method calls.
"""
encoded_image = image_utils.encode_image(image_matrix, image_format)
height, width, num_channels = image_matrix.shape
return self.add_encoded_image_feature(encoded_image, image_format, height,
width, num_channels, image_source_id,
feature_prefix, label)
def add_encoded_image_feature(
self,
encoded_image: bytes,
image_format: Optional[str] = None,
height: Optional[int] = None,
width: Optional[int] = None,
num_channels: Optional[int] = None,
image_source_id: Optional[bytes] = None,
feature_prefix: Optional[str] = None,
label: Optional[Union[int, Sequence[int]]] = None) -> 'TfExampleBuilder':
"""Adds encoded image features to the example.
See `tf_example_feature_key.EncodedImageFeatureKey` for list of feature keys
that will be added to the example.
Image format, height, width, and channels are inferred from the encoded
image bytes if any of them is not provided. Hashed image will be used if
pre-generated source ID is not provided.
Example usages:
>>> example_builder = TfExampleBuilder()
* For adding RGB image feature:
>>> example_builder.add_encoded_image_feature(image_bytes)
* For adding RGB image feature with pre-generated source ID:
>>> example_builder.add_encoded_image_feature(
image_bytes, image_source_id=image_source_id)
* For adding single-channel depth image feature:
>>> example_builder.add_encoded_image_feature(
image_bytes, feature_prefix='depth')
Args:
encoded_image: Encoded image string.
image_format: Image format string.
height: Number of rows.
width: Number of columns.
num_channels: Number of channels.
image_source_id: Unique string ID to identify the image.
feature_prefix: Feature prefix for image features.
label: the label or a list of labels for the image.
Returns:
The builder object for subsequent method calls.
"""
if image_format == 'RAW':
if not (height and width and num_channels):
raise ValueError('For raw image feature, height, width and '
'num_channels fields are required.')
if not all((height, width, num_channels, image_format)):
(height, width, num_channels, image_format) = (
image_utils.decode_image_metadata(encoded_image))
else:
image_format = image_utils.validate_image_format(image_format)
feature_key = tf_example_feature_key.EncodedImageFeatureKey(feature_prefix)
# If source ID is not provided, we use hashed encoded image as the source
# ID. Note that we only keep 24 bits to be consistent with the Model Garden
# requirement, which will transform the source ID into float32.
if not image_source_id:
hashed_image = int(hashlib.blake2s(encoded_image).hexdigest(), 16)
image_source_id = _to_bytes(str(hashed_image % ((1 << 24) + 1)))
if label is not None:
self.add_ints_feature(feature_key.label, label)
return (
self.add_bytes_feature(feature_key.encoded, encoded_image)
.add_bytes_feature(feature_key.format, image_format)
.add_ints_feature(feature_key.height, [height])
.add_ints_feature(feature_key.width, [width])
.add_ints_feature(feature_key.num_channels, num_channels)
.add_bytes_feature(feature_key.source_id, image_source_id))
def add_boxes_feature(
self,
xmins: Sequence[float],
xmaxs: Sequence[float],
ymins: Sequence[float],
ymaxs: Sequence[float],
labels: Sequence[int],
confidences: Optional[Sequence[float]] = None,
normalized: bool = True,
feature_prefix: Optional[str] = None) -> 'TfExampleBuilder':
"""Adds box and label features to the example.
Four features will be generated for xmin, ymin, xmax, and ymax. One feature
will be generated for label. Different feature keys will be used for
normalized boxes and pixel-value boxes, depending on the value of
`normalized`.
Example usages:
>>> example_builder = TfExampleBuilder()
>>> example_builder.add_boxes_feature(xmins, xmaxs, ymins, ymaxs, labels)
Args:
xmins: A list of minimum X coordinates.
xmaxs: A list of maximum X coordinates.
ymins: A list of minimum Y coordinates.
ymaxs: A list of maximum Y coordinates.
labels: The labels of added boxes.
confidences: The confidences of added boxes.
normalized: Indicate if the coordinates of boxes are normalized.
feature_prefix: Feature prefix for added box features.
Returns:
The builder object for subsequent method calls.
"""
if normalized:
feature_key = tf_example_feature_key.BoxFeatureKey(feature_prefix)
else:
feature_key = tf_example_feature_key.BoxPixelFeatureKey(feature_prefix)
self.add_floats_feature(feature_key.xmin, xmins)
self.add_floats_feature(feature_key.xmax, xmaxs)
self.add_floats_feature(feature_key.ymin, ymins)
self.add_floats_feature(feature_key.ymax, ymaxs)
self.add_ints_feature(feature_key.label, labels)
if confidences is not None:
self.add_floats_feature(feature_key.confidence, confidences)
return self
def _compute_mask_areas(
self, instance_mask_matrices: np.ndarray) -> Sequence[float]:
return np.sum(
instance_mask_matrices, axis=(1, 2, 3),
dtype=float).flatten().tolist()
def add_instance_mask_matrices_feature(
self,
instance_mask_matrices: np.ndarray,
feature_prefix: Optional[str] = None) -> 'TfExampleBuilder':
"""Encodes and adds instance mask features to the example.
See `tf_example_feature_key.EncodedInstanceMaskFeatureKey` for list of
feature keys that will be added to the example. Please note that all masks
will be encoded as PNG images.
Example usages:
>>> example_builder = TfExampleBuilder()
>>> example_builder.add_instance_mask_matrices_feature(
instance_mask_matrices)
TODO(b/223653024): Provide a way to generate visualization mask from
feature mask.
Args:
instance_mask_matrices: Numpy instance mask matrices with shape
(num_instance, height, width, 1) or (num_instance, height, width).
feature_prefix: Feature prefix for instance mask features.
Returns:
The builder object for subsequent method calls.
"""
if len(instance_mask_matrices.shape) == 3:
instance_mask_matrices = instance_mask_matrices[..., np.newaxis]
mask_areas = self._compute_mask_areas(instance_mask_matrices)
encoded_instance_masks = list(
map(lambda x: image_utils.encode_image(x, 'PNG'),
instance_mask_matrices))
return self.add_encoded_instance_masks_feature(encoded_instance_masks,
mask_areas, feature_prefix)
def add_encoded_instance_masks_feature(
self,
encoded_instance_masks: Sequence[bytes],
mask_areas: Optional[Sequence[float]] = None,
feature_prefix: Optional[str] = None) -> 'TfExampleBuilder':
"""Adds encoded instance mask features to the example.
See `tf_example_feature_key.EncodedInstanceMaskFeatureKey` for list of
feature keys that will be added to the example.
Image area is inferred from the encoded instance mask bytes if not provided.
Example usages:
>>> example_builder = TfExampleBuilder()
>>> example_builder.add_encoded_instance_masks_feature(
instance_mask_bytes)
TODO(b/223653024): Provide a way to generate visualization mask from
feature mask.
Args:
encoded_instance_masks: A list of encoded instance mask string. Note that
the encoding is not changed in this function and it always assumes the
image is in "PNG" format.
mask_areas: Areas for each instance masks.
feature_prefix: Feature prefix for instance mask features.
Returns:
The builder object for subsequent method calls.
"""
encoded_instance_masks = _to_bytes_array(encoded_instance_masks)
if mask_areas is None:
instance_mask_matrices = np.array(
list(map(image_utils.decode_image, encoded_instance_masks)))
mask_areas = self._compute_mask_areas(instance_mask_matrices)
feature_key = tf_example_feature_key.EncodedInstanceMaskFeatureKey(
feature_prefix)
return (
self.add_bytes_feature(feature_key.mask, encoded_instance_masks)
.add_floats_feature(feature_key.area, mask_areas))
def add_semantic_mask_matrix_feature(
self,
mask_matrix: np.ndarray,
mask_format: str = 'PNG',
visualization_mask_matrix: Optional[np.ndarray] = None,
visualization_mask_format: str = 'PNG',
feature_prefix: Optional[str] = None) -> 'TfExampleBuilder':
"""Encodes and adds semantic mask features to the example.
See `tf_example_feature_key.EncodedSemanticMaskFeatureKey` for list of
feature keys that will be added to the example.
Example usages:
>>> example_builder = TfExampleBuilder()
* For adding semantic mask feature:
>>> example_builder.add_semantic_mask_matrix_feature(
semantic_mask_matrix)
* For adding semantic mask feature and visualization mask feature:
>>> example_builder.add_semantic_mask_matrix_feature(
semantic_mask_matrix,
visualization_mask_matrix=visualization_mask_matrix)
* For adding predicted semantic mask feature with visualization mask:
>>> example_builder.add_encoded_semantic_mask_feature(
predicted_mask_matrix,
visualization_mask_matrix=predicted_visualization_mask_matrix,
feature_prefix='predicted')
TODO(b/223653024): Provide a way to generate visualization mask from
feature mask.
Args:
mask_matrix: Numpy semantic mask matrix with shape (height, width, 1) or
(height, width).
mask_format: Mask format string, defaults to 'PNG'.
visualization_mask_matrix: Numpy visualization mask matrix for semantic
mask with shape (height, width, 3).
visualization_mask_format: Visualization mask format string, defaults to
'PNG'.
feature_prefix: Feature prefix for semantic mask features.
Returns:
The builder object for subsequent method calls.
"""
if len(mask_matrix.shape) == 2:
mask_matrix = mask_matrix[..., np.newaxis]
encoded_mask = image_utils.encode_image(mask_matrix, mask_format)
encoded_visualization_mask = None
if visualization_mask_matrix is not None:
encoded_visualization_mask = image_utils.encode_image(
visualization_mask_matrix, visualization_mask_format)
return self.add_encoded_semantic_mask_feature(encoded_mask, mask_format,
encoded_visualization_mask,
visualization_mask_format,
feature_prefix)
def add_encoded_semantic_mask_feature(
self, encoded_mask: bytes,
mask_format: str = 'PNG',
encoded_visualization_mask: Optional[bytes] = None,
visualization_mask_format: str = 'PNG',
feature_prefix: Optional[str] = None) -> 'TfExampleBuilder':
"""Adds encoded semantic mask features to the example.
See `tf_example_feature_key.EncodedSemanticMaskFeatureKey` for list of
feature keys that will be added to the example.
Example usages:
>>> example_builder = TfExampleBuilder()
* For adding semantic mask feature:
>>> example_builder.add_encoded_semantic_mask_feature(semantic_mask_bytes)
* For adding semantic mask feature and visualization mask feature:
>>> example_builder.add_encoded_semantic_mask_feature(
semantic_mask_bytes,
encoded_visualization_mask=visualization_mask_bytes)
* For adding predicted semantic mask feature with visualization mask:
>>> example_builder.add_encoded_semantic_mask_feature(
predicted_mask_bytes,
encoded_visualization_mask=predicted_visualization_mask_bytes,
feature_prefix='predicted')
TODO(b/223653024): Provide a way to generate visualization mask from
feature mask.
Args:
encoded_mask: Encoded semantic mask string.
mask_format: Semantic mask format string, defaults to 'PNG'.
encoded_visualization_mask: Encoded visualization mask string.
visualization_mask_format: Visualization mask format string, defaults to
'PNG'.
feature_prefix: Feature prefix for semantic mask features.
Returns:
The builder object for subsequent method calls.
"""
feature_key = tf_example_feature_key.EncodedSemanticMaskFeatureKey(
feature_prefix)
example_builder = (
self.add_bytes_feature(feature_key.mask, encoded_mask)
.add_bytes_feature(feature_key.mask_format, mask_format))
if encoded_visualization_mask is not None:
example_builder = (
example_builder.add_bytes_feature(
feature_key.visualization_mask, encoded_visualization_mask)
.add_bytes_feature(
feature_key.visualization_mask_format, visualization_mask_format))
return example_builder
def add_panoptic_mask_matrix_feature(
self,
panoptic_category_mask_matrix: np.ndarray,
panoptic_instance_mask_matrix: np.ndarray,
panoptic_category_mask_format: str = 'PNG',
panoptic_instance_mask_format: str = 'PNG',
feature_prefix: Optional[str] = None) -> 'TfExampleBuilder':
"""Encodes and adds panoptic mask features to the example.
See `tf_example_feature_key.EncodedPanopticMaskFeatureKey` for list of
feature keys that will be added to the example.
Example usages:
>>> example_builder = TfExampleBuilder()
>>> example_builder.add_panoptic_mask_matrix_feature(
panoptic_category_mask_matrix, panoptic_instance_mask_matrix)
TODO(b/223653024): Provide a way to generate visualization mask from
feature mask.
Args:
panoptic_category_mask_matrix: Numpy panoptic category mask matrix with
shape (height, width, 1) or (height, width).
panoptic_instance_mask_matrix: Numpy panoptic instance mask matrix with
shape (height, width, 1) or (height, width).
panoptic_category_mask_format: Panoptic category mask format string,
defaults to 'PNG'.
panoptic_instance_mask_format: Panoptic instance mask format string,
defaults to 'PNG'.
feature_prefix: Feature prefix for panoptic mask features.
Returns:
The builder object for subsequent method calls.
"""
if len(panoptic_category_mask_matrix.shape) == 2:
panoptic_category_mask_matrix = (
panoptic_category_mask_matrix[..., np.newaxis])
if len(panoptic_instance_mask_matrix.shape) == 2:
panoptic_instance_mask_matrix = (
panoptic_instance_mask_matrix[..., np.newaxis])
encoded_panoptic_category_mask = image_utils.encode_image(
panoptic_category_mask_matrix, panoptic_category_mask_format)
encoded_panoptic_instance_mask = image_utils.encode_image(
panoptic_instance_mask_matrix, panoptic_instance_mask_format)
return self.add_encoded_panoptic_mask_feature(
encoded_panoptic_category_mask, encoded_panoptic_instance_mask,
panoptic_category_mask_format, panoptic_instance_mask_format,
feature_prefix)
def add_encoded_panoptic_mask_feature(
self,
encoded_panoptic_category_mask: bytes,
encoded_panoptic_instance_mask: bytes,
panoptic_category_mask_format: str = 'PNG',
panoptic_instance_mask_format: str = 'PNG',
feature_prefix: Optional[str] = None) -> 'TfExampleBuilder':
"""Adds encoded panoptic mask features to the example.
See `tf_example_feature_key.EncodedPanopticMaskFeatureKey` for list of
feature keys that will be added to the example.
Example usages:
>>> example_builder = TfExampleBuilder()
>>> example_builder.add_encoded_panoptic_mask_feature(
encoded_panoptic_category_mask, encoded_panoptic_instance_mask)
TODO(b/223653024): Provide a way to generate visualization mask from
feature mask.
Args:
encoded_panoptic_category_mask: Encoded panoptic category mask string.
encoded_panoptic_instance_mask: Encoded panoptic instance mask string.
panoptic_category_mask_format: Panoptic category mask format string,
defaults to 'PNG'.
panoptic_instance_mask_format: Panoptic instance mask format string,
defaults to 'PNG'.
feature_prefix: Feature prefix for panoptic mask features.
Returns:
The builder object for subsequent method calls.
"""
feature_key = tf_example_feature_key.EncodedPanopticMaskFeatureKey(
feature_prefix)
return (
self.add_bytes_feature(
feature_key.category_mask, encoded_panoptic_category_mask)
.add_bytes_feature(
feature_key.category_mask_format, panoptic_category_mask_format)
.add_bytes_feature(
feature_key.instance_mask, encoded_panoptic_instance_mask)
.add_bytes_feature(
feature_key.instance_mask_format, panoptic_instance_mask_format))