# AUTOGENERATED! DO NOT EDIT! File to edit: tensorflow-tfrecord_builder.ipynb (unless otherwise specified).

__all__ = ['logger', 'create_tfrecord_entry', 'create_tfrecord_file', 'int64_feature', 'int64_list_feature',
           'bytes_feature', 'bytes_list_feature', 'float_list_feature', 'create_labelmap_file']


# Cell
import logging
import io
import tensorflow as tf
from os import environ
from os.path import basename
from ..image.pillow_tools import get_image_size


# Cell
environ['TF_CPP_MIN_LOG_LEVEL'] = '2'    # Suppress TensorFlow logging (1)

logger = logging.getLogger(__name__)


# Cell
def create_tfrecord_entry(categories, annotation):
    """
    Create tfrecord entry with annotations for one file / image.
    `categories`: the categories used
    `annotation`: the annotation of a file / image
    return: the tfrecord entry
    """
    with tf.io.gfile.GFile(annotation.file_path, 'rb') as fid:
        encoded_jpg = fid.read()
    _, width, height = get_image_size(io.BytesIO(encoded_jpg))

    file_name = basename(annotation.file_path).encode('utf8')
    image_format = b'jpg'
    xmins = []
    xmaxs = []
    ymins = []
    ymaxs = []
    classes_text = []
    classes = []

    for region in annotation.regions:
        x_min = min(region.points_x)
        y_min = min(region.points_y)
        x_max = max(region.points_x)
        y_max = max(region.points_y)
        category = region.labels[0] if len(region.labels) else ''

        xmins.append(x_min / width)
        xmaxs.append(x_max / width)
        ymins.append(y_min / height)
        ymaxs.append(y_max / height)
        classes_text.append(category.encode('utf8'))
        classes.append(categories.index(category))

    tf_example = tf.train.Example(features=tf.train.Features(feature={
        'image/height': int64_feature(height),
        'image/width': int64_feature(width),
        'image/filename': bytes_feature(file_name),
        'image/source_id': bytes_feature(file_name),
        'image/encoded': bytes_feature(encoded_jpg),
        'image/format': bytes_feature(image_format),
        'image/object/bbox/xmin': float_list_feature(xmins),
        'image/object/bbox/xmax': float_list_feature(xmaxs),
        'image/object/bbox/ymin': float_list_feature(ymins),
        'image/object/bbox/ymax': float_list_feature(ymaxs),
        'image/object/class/text': bytes_list_feature(classes_text),
        'image/object/class/label': int64_list_feature(classes),
    }))
    return tf_example


# Cell
def create_tfrecord_file(output_path, categories, annotations):
    """
    Create a tfrecord file for a sub-data-set, which can be one of the following: training, validation, test
    `output_path`: the path including the filename of the tfrecord file
    `categories`: the categories used
    `annotations`: the annotations of the files / images
    """
    writer = tf.io.TFRecordWriter(output_path)
    for annotation in annotations.values():
        tf_example = create_tfrecord_entry(categories, annotation)
        writer.write(tf_example.SerializeToString())
    writer.close()
    logger.info('Successfully created the TFRecord file: {}'.format(output_path))


# Cell
def int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def int64_list_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))


def bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def bytes_list_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))


def float_list_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))


# Cell
def create_labelmap_file(output_path, categories, start=1):
    """
    Create labelmap protobuffer text file containing the categories.
    Format is compatible with Tensorflow Object Detection API.
    For object detection data-sets, the categories should exclude the background class and `start` should be 1.
    `output_path`: the path including the filename of the protobuffer text file
    `categories`: a list of the categories to write
    `start`: the category index for the first category
    """
    # create label_map data
    label_map = ''
    for index, category in enumerate(categories, start=start):
        label_map = label_map + "item {\n"
        label_map = label_map + " id: " + str(index) + "\n"
        label_map = label_map + " name: '" + category + "'\n}\n\n"
    label_map = label_map[:-1]

    # write label_map file
    with open(output_path, "w") as f:
        f.write(label_map)
        f.close()
