static_dataloader.py

###############################################################################
# Copyright (C) 2020-2021 Habana Labs, Ltd. an Intel Company
###############################################################################

import tensorflow as tf
from TensorFlow.computer_vision.SSD_ResNet34.constants import *


def deserialize(data, batch_size, dtype):
    deserialized = tf.io.parse_tensor(data, tf.string)
    image, num_matched_boxes, boxes, classes = tf.unstack(deserialized, 4)

    image = tf.io.parse_tensor(image, tf.float32)
    num_matched_boxes = tf.io.parse_tensor(num_matched_boxes, tf.float32)
    boxes = tf.io.parse_tensor(boxes, tf.float32)
    classes = tf.io.parse_tensor(classes, tf.int32)

    if dtype != tf.float32:
        image = tf.cast(image, dtype)

    IMAGE_CHANNELS = 3
    BOXES_COORDS = 4
    image = tf.reshape(image, (batch_size, IMAGE_SIZE,
                               IMAGE_SIZE, IMAGE_CHANNELS))
    num_matched_boxes = tf.reshape(num_matched_boxes, (batch_size,))
    boxes = tf.reshape(boxes, (batch_size, NUM_SSD_BOXES, BOXES_COORDS))
    classes = tf.reshape(classes, (batch_size, NUM_SSD_BOXES))

    features = dict(num_matched_boxes=num_matched_boxes,
                    boxes=boxes, classes=classes)
    return image, features


class SSDStaticInputReader(object):
    """Static input reader for dataset.
    Reads data generated by generate_static_dataset.py"""

    def __init__(self,
                 file_pattern):
        self.file_pattern = file_pattern

    def __call__(self, params):
        dtype_conversion = dict(bf16=tf.bfloat16, fp32=tf.float32)

        batch_size = params['batch_size']
        dtype = dtype_conversion[params['dtype']]

        dataset = tf.data.Dataset.list_files(
            self.file_pattern, shuffle=False)
        dataset = tf.data.TFRecordDataset(dataset)
        def convert(data): return deserialize(data, batch_size, dtype)
        dataset = dataset.map(convert)
        dataset = dataset.repeat()
        dataset = dataset.prefetch(8)

        return dataset