from_tfrecords.py

import tensorflow as tf


__all__ = ['from_tfrecords']


def from_tfrecords(file_paths, compression_type=None, features=None, feature_lists=None):
    dataset = tf.data.TFRecordDataset(file_paths, compression_type=compression_type)

    if not features and not feature_lists:
        features, feature_lists = detect_schema(dataset)

    if feature_lists:
        parser = read_sequence_example(features, feature_lists)
    else:
        parser = read_example(features)

    parsed = dataset.map(parser)
    return (i for i in parsed)


def read_example(features):

    def parse(serialized):
        example = tf.io.parse_single_example(
            serialized,
            features=features)
        return example
    
    return parse


def read_sequence_example(features, feature_lists):

    def parse(serialized):
        context, sequence = tf.io.parse_single_sequence_example(
            serialized,
            context_features=features,
            sequence_features=feature_lists)
        context.update(sequence)
        return context
    
    return parse


def detect_schema(dataset):
    features = None
    feature_lists = None

    serialized = next(iter(dataset.map(lambda serialized: serialized)))
    seq_ex = tf.train.SequenceExample.FromString(serialized.numpy())


    if seq_ex.context.feature:
        features = {}
        for key, feature in seq_ex.context.feature.items():
            features[key] = tf.io.FixedLenFeature((), get_type(feature))

    if seq_ex.feature_lists.feature_list:
        feature_lists = {}
        for key, feature_list in seq_ex.feature_lists.feature_list.items():
            feature_lists[key] = tf.io.FixedLenSequenceFeature((), get_type(feature_list.feature[0]))

    return features, feature_lists


def get_type(feature):
    if feature.HasField('int64_list'):
        return tf.int64
    if feature.HasField('float_list'):
        return tf.float32
    if feature.HasField('bytes_list'):
        return tf.string