# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Bosch Dataset Preperation

import tensorflow as tf
from object_detection.utils import dataset_util
import yaml
import os, sys
from PIL import Image

# ## Define Labels

LABEL_DICT = {
    "Green" : 1,
    "Red" : 2,
    "GreenLeft" : 3,
    "GreenRight" : 4,
    "RedLeft" : 5,
    "RedRight" : 6,
    "Yellow" : 7,
    "off" : 8,
    "RedStraight" : 9,
    "GreenStraight" : 10,
    "GreenStraightLeft" : 11,
    "GreenStraightRight" : 12,
    "RedStraightLeft" : 13,
    "RedStraightRight" : 14
}


# ## Process each Record

def create_tf_record(sample):
    
    # Image dimensions
    img_height = 720
    img_width = 1280

    fname = sample['path'] 
    fname = fname.encode()

    with tf.gfile.GFile(sample['path'], 'rb') as fid:
        encoded_image = fid.read()

    image_format = 'png'.encode() 

    # Bounding box location lists for the given record
    xmins = []
    xmaxs = []
    ymins = []
    ymaxs = []
    
    # Class descriptions and classes lists
    classes_text = [] 
    classes = []

    for box in sample['boxes']:

        xmins.append(float(box['x_min'] / img_width))
        xmaxs.append(float(box['x_max'] / img_width))
        ymins.append(float(box['y_min'] / img_height))
        ymaxs.append(float(box['y_max'] / img_height))
        classes_text.append(box['label'].encode())
        classes.append(int(LABEL_DICT[box['label']]))


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

    return tf_record


# ##  Process all records and create tf record for training and evaluation dataset

# +
def main(_):

    train_writer = tf.python_io.TFRecordWriter('/data2/TrafficLight/tf_record/train.tf_record')
    eval_writer = tf.python_io.TFRecordWriter('/data2/TrafficLight/tf_record/eval.tf_record')


    INPUT_YAML = '/data2/TrafficLight/train.yaml'
    examples = yaml.load(open(INPUT_YAML, 'rb').read())

    len_examples = len(examples)
    print("Loaded ", len(examples), "examples")

    for i in range(len(examples)):
        examples[i]['path'] = os.path.abspath(os.path.join(os.path.dirname(INPUT_YAML), examples[i]['path']))
        

    # Generate Training TF records

    counter = 0
    for example in examples[:3800]:
        tf_example = create_tf_record(example)
        train_writer.write(tf_example.SerializeToString())

        if counter % 20 == 0:
            print("Percent done", (float(counter)/float(len(examples[:3800])))*100)
        counter += 1

    #train_writer.close()
    
    # Generate Evaluation TF records
    
    counter = 0
    for example in examples[3801:4600]:
        tf_example = create_tf_record(example)
        eval_writer.write(tf_example.SerializeToString())

        if counter % 20 == 0:
            print("Percent done", (float(counter)/float(len(examples[3801:4600])))*100)
        counter += 1
    
    #eval_writer.close()


if __name__ == '__main__':
    tf.app.run()
# -



StarcoderdataJupyterScriptsDedupFiltered