def create_tf_example(image: np.ndarray,
window: Box,
labels: np.ndarray,
class_map: ClassMap,
chip_id: str = ''):
"""Create a TensorFlow Example from an image, the labels, &c.
Args:
image: An np.ndarray containing the image data.
window: A Box object containing the bounding box for this example.
labels: An nd.array containing the label data.
class_map: A ClassMap object containing mappings between
numerical and textual labels.
chip_id: The chip id as a string.
Returns:
A DeepLab-compatible TensorFlow Example object containing the
given data.
"""
import tensorflow as tf
from object_detection.utils import dataset_util
class_keys = set(class_map.get_keys())
def _clean(n):
return (n if n in class_keys else 0x00)
clean = np.vectorize(_clean, otypes=[np.uint8])
image_encoded = numpy_to_png(image)
image_filename = chip_id.encode('utf8')
image_format = 'png'.encode('utf8')
image_height, image_width, image_channels = image.shape
image_segmentation_class_encoded = numpy_to_png(clean(labels))
image_segmentation_class_format = 'png'.encode('utf8')
features = tf.train.Features(
feature={
'image/encoded':
dataset_util.bytes_feature(image_encoded),
'image/filename':
dataset_util.bytes_feature(image_filename),
'image/format':
dataset_util.bytes_feature(image_format),
'image/height':
dataset_util.int64_feature(image_height),
'image/width':
dataset_util.int64_feature(image_width),
'image/channels':
dataset_util.int64_feature(image_channels),
'image/segmentation/class/encoded':
dataset_util.bytes_feature(image_segmentation_class_encoded),
'image/segmentation/class/format':
dataset_util.bytes_feature(image_segmentation_class_format),
})
return tf.train.Example(features=features)
def make_debug_images(record_path: str, output_dir: str, class_map: ClassMap,
p: float) -> None:
"""Render a random sample of the TFRecords in a given file as
human-viewable PNG files.
Args:
record_path: Path to the TFRecord file.
output_dir: Destination directory for the generated PNG files.
p: The probability of rendering a particular record.
Returns:
None
"""
import tensorflow as tf
make_dir(output_dir)
ids = class_map.get_keys()
color_strs = list(map(lambda c: c.color, class_map.get_items()))
color_ints = list(map(lambda c: color_to_integer(c), color_strs))
correspondence = dict(zip(ids, color_ints))
def _label_fn(v: int) -> int:
if v in correspondence:
return correspondence.get(v)
else:
return 0
label_fn = np.vectorize(_label_fn, otypes=[np.uint64])
def _image_fn(pixel: int) -> int:
if (pixel & 0x00ffffff):
r = ((pixel >> 41 & 0x7f) + (pixel >> 17 & 0x7f)) << 16
g = ((pixel >> 33 & 0x7f) + (pixel >> 9 & 0x7f)) << 8
b = ((pixel >> 25 & 0x7f) + (pixel >> 1 & 0x7f)) << 0
return r + g + b
else:
return pixel >> 24
image_fn = np.vectorize(_image_fn, otypes=[np.uint64])
log.info('Generating debug chips')
tfrecord_iter = tf.python_io.tf_record_iterator(record_path)
for ind, example in enumerate(tfrecord_iter):
if np.random.rand() <= p:
example = tf.train.Example.FromString(example)
im_unpacked, labels = parse_tf_example(example)
im_r = np.array(im_unpacked[:, :, 0], dtype=np.uint64) * 1 << 40
im_g = np.array(im_unpacked[:, :, 1], dtype=np.uint64) * 1 << 32
im_b = np.array(im_unpacked[:, :, 2], dtype=np.uint64) * 1 << 24
im_packed = im_r + im_g + im_b
labels_packed = label_fn(np.array(labels))
im_labels_packed = im_packed + labels_packed
im_packed = image_fn(im_labels_packed)
im_unpacked[:, :, 0] = np.bitwise_and(
im_packed >> 16, 0xff, dtype=np.uint8)
im_unpacked[:, :, 1] = np.bitwise_and(
im_packed >> 8, 0xff, dtype=np.uint8)
im_unpacked[:, :, 2] = np.bitwise_and(
im_packed >> 0, 0xff, dtype=np.uint8)
output_path = join(output_dir, '{}.png'.format(ind))
save_img(im_unpacked, output_path)