# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Utility functions for creating TFRecord data sets.
source: https://github.com/tensorflow/models/blob/master/research/object_detection/utils/dataset_util.py
"""

import json
import tensorflow as tf
from multiprocessing.pool import ThreadPool
from threading import Lock
from utils.preprocessing import decode_shi_labels
import numpy as np
from time import time, sleep
from PIL import Image
import sys
import os
import cv2
from utils.base_util import is_image, random_select, random_integers, Timer, file_list
from matplotlib import pyplot as plt
from random import shuffle


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))


def read_examples_list(path):
    """Read list of training or validation examples.

    The file is assumed to contain a single example per line where the first
    token in the line is an identifier that allows us to find the image and
    annotation xml for that example.

    For example, the line:
    xyz 3
    would allow us to find files xyz.jpg and xyz.xml (the 3 would be ignored).

    Args:
      path: absolute path to examples list file.

    Returns:
      list of example identifiers (strings).
    """
    with tf.gfile.GFile(path) as fid:
        lines = fid.readlines()
    return [line.strip().split(' ')[0] for line in lines]


def recursive_parse_xml_to_dict(xml):
    """Recursively parses XML contents to python dict.

    We assume that `object` tags are the only ones that can appear
    multiple times at the same level of a tree.

    Args:
      xml: xml tree obtained by parsing XML file contents using lxml.etree

    Returns:
      Python dictionary holding XML contents.
    """
    if not xml:
        return {xml.tag: xml.text}
    result = {}
    for child in xml:
        child_result = recursive_parse_xml_to_dict(child)
        if child.tag != 'object':
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}


def make_initializable_iterator(dataset):
    """Creates an iterator, and initializes tables.

    This is useful in cases where make_one_shot_iterator wouldn't work because
    the graph contains a hash table that needs to be initialized.

    Args:
      dataset: A `tf.data.Dataset` object.

    Returns:
      A `tf.data.Iterator`.
    """
    iterator = dataset.make_initializable_iterator()
    tf.add_to_collection(tf.GraphKeys.TABLE_INITIALIZERS, iterator.initializer)
    return iterator


def read_dataset(
        file_read_func, decode_func, input_files, config, num_workers=1,
        worker_index=0):
    """Reads a dataset, and handles repetition and shuffling.

    Args:
      file_read_func: Function to use in tf.data.Dataset.interleave, to read
        every individual file into a tf.data.Dataset.
      decode_func: Function to apply to all records.
      input_files: A list of file paths to read.
      config: A input_reader_builder.InputReader object.
      num_workers: Number of workers / shards.
      worker_index: Id for the current worker.

    Returns:
      A tf.data.Dataset based on config.
    """
    # Shard, shuffle, and read files.
    filenames = tf.concat([tf.matching_files(pattern) for pattern in input_files],
                          0)
    dataset = tf.data.Dataset.from_tensor_slices(filenames)
    dataset = dataset.shard(num_workers, worker_index)
    dataset = dataset.repeat(config.num_epochs or None)
    if config.shuffle:
        dataset = dataset.shuffle(config.filenames_shuffle_buffer_size,
                                  reshuffle_each_iteration=True)

    # Read file records and shuffle them.
    # If cycle_length is larger than the number of files, more than one reader
    # will be assigned to the same file, leading to repetition.
    cycle_length = tf.cast(
        tf.minimum(config.num_readers, tf.size(filenames)), tf.int64)
    # TODO: find the optimal block_length.
    dataset = dataset.interleave(
        file_read_func, cycle_length=cycle_length, block_length=1)

    if config.shuffle:
        dataset = dataset.shuffle(config.shuffle_buffer_size,
                                  reshuffle_each_iteration=True)

    dataset = dataset.map(decode_func, num_parallel_calls=config.num_readers)
    return dataset.prefetch(config.prefetch_buffer_size)


def run_thread_blocks(number_of_threads, data_list, thread_fn, writer=None, timer=None, return_result=False):
    pool = ThreadPool(processes=number_of_threads)
    data_block_list = list()
    partition_size = int(len(data_list) / number_of_threads)
    for i in range(number_of_threads):
        if i != (number_of_threads - 1):
            data_block_list.append(data_list[i * partition_size:(i + 1) * partition_size])
        else:
            data_block_list.append(data_list[i * partition_size:])

    threads = list()
    for i in range(number_of_threads):
        threads.append(pool.apply_async(thread_fn, (i, data_block_list[i], writer, timer)))

    if return_result:
        outputs = list()
        for i in range(number_of_threads):
            outputs.append(threads[i].get())
        return outputs
    else:
        for i in range(number_of_threads):
            threads[i].get()


label_mask = [
    [0, 0, 255],  # Sea
    [0, 255, 0],  # Sky
    [255, 255, 255],  # Ship
    [255, 0, 0],  # Shore
    [255, 255, 0],  # Own ship
    [100, 100, 100],  # Crane
    [255, 0, 255],  # Etc
    [0, 0, 0]  # Ignore
]
label_mask = np.array(label_mask).astype(np.float32)


def convert_label(label_path, image_size, apply_ignore=True):
    """
    Convert RGB label image to onehot label image
    :param label_path: File path of RGB label image
    :param image_size: Size to resize result image
    :param apply_ignore: Apply ignore
    :return:
    """
    label = np.array(Image.open(label_path).resize((image_size[0], image_size[1]), Image.ANTIALIAS))[:, :, :3]
    converted_label = np.zeros(shape=[image_size[0], image_size[1]], dtype=np.float32)
    for index, mask in enumerate(label_mask):
        selected = np.all(label == mask, axis=2, keepdims=False)
        converted_label += selected.astype(np.uint8) * index

    ignored_part = np.equal(converted_label, len(label_mask) - 1).astype(np.float32)
    if apply_ignore:
        converted_label += ignored_part * (255 - len(label_mask) + 1)

    return converted_label


def convert_to_nearest_label(label_path, image_size, apply_ignore=True):
    """
    Convert RGB label image to onehot label image
    :param label_path: File path of RGB label image
    :param image_size: Size to resize result image
    :param apply_ignore: Apply ignore
    :return:
    """
    label = np.array(Image.open(label_path).resize((image_size[0], image_size[1]), Image.ANTIALIAS))[:, :, :3]
    label = label.astype(np.float32)
    stacked_label = list()
    for index, mask in enumerate(label_mask):
        length = np.sum(cv2.pow(label - mask, 2), axis=2, keepdims=False)
        length = cv2.sqrt(length)
        stacked_label.append(length)

    stacked_label = np.array(stacked_label)
    stacked_label = np.transpose(stacked_label, [1, 2, 0])
    converted_to_classes = np.argmin(stacked_label, axis=2).astype(np.uint8)
    if apply_ignore:
        ignore_mask = (converted_to_classes == (len(label_mask) - 1)).astype(np.uint8)
        ignore_mask *= (256 - len(label_mask))
        converted_to_classes += ignore_mask

    return converted_to_classes


def make_shi_tfrecord(image_path, label_path, train_data_output_path, eval_data_output_path, shuffle_data=True):
    """
    Multi thread version of tfrecord maker
    :param image_path: Path of images
    :param label_path: Path of image dataset
    :param train_data_output_path: Path of train tfrecord
    :param eval_data_output_path: Path of eval tfrecord
    :param shuffle_data: To shuffle data or not
    :return:
    """
    if isinstance(image_path, str):
        image_folder_list = [os.path.join(image_path, l) for l in os.listdir(image_path)
                             if is_image(os.path.join(image_path, l))]
    else:
        raise ValueError('Need to provide path of image set in string or list of images')

    if isinstance(label_path, str):
        label_folder_list = [os.path.join(label_path, l) for l in os.listdir(label_path)
                             if is_image(os.path.join(label_path, l))]
    else:
        raise ValueError('Need to provide path of image set in string or list of images')
    print('Make tfrecord from %s to %s' % (label_path, train_data_output_path))

    image_and_label_list = [[i, j] for [i, j] in zip(image_folder_list, label_folder_list)]
    if shuffle_data:
        shuffle(image_and_label_list)

    eval_list, train_list = random_select(image_and_label_list, 200, return_unselected=True)

    train_writer = tf.python_io.TFRecordWriter(train_data_output_path)
    eval_writer = tf.python_io.TFRecordWriter(eval_data_output_path)

    # Make list of image data
    print('Start make image and label pairs')

    mutex = Lock()
    mutex2 = Lock()

    image_size = (2048, 2048)
    division = 1

    def load_image_pairs(thread_id, img_and_label_list, writer=None, timer=None):
        global finished_samples
        global time_mean
        global time_start
        if thread_id == 0:
            time_start = time()
            finished_samples = 0
            time_mean = -1
        _height = int(image_size[0] / division)
        _width = int(image_size[1] / division)
        for idx, (img, lab) in enumerate(img_and_label_list):
            image = np.array(Image.open(img).convert('RGB').resize((image_size[0], image_size[1]), Image.ANTIALIAS))
            converted_label = convert_to_nearest_label(lab, image_size)

            divided_images = list()
            divided_labels = list()
            for d_i in range(division * 2 - 1):
                for d_j in range(division * 2 - 1):
                    selected_label = converted_label[
                        d_i * int(0.5 * _height):d_i * int(0.5 * _height) + _height,
                        d_j * int(0.5 * _width):d_j * int(0.5 * _width) + _width
                    ]

                    if np.sum(np.equal(selected_label, 2).astype(np.int32)) == 0:
                        continue

                    divided_images.append(
                        np.array(
                            image[
                                d_i * int(0.5 * _height):d_i * int(0.5 * _height) + _height,
                                d_j * int(0.5 * _width):d_j * int(0.5 * _width) + _width,
                                :3
                            ],
                            dtype=np.uint8
                        )
                    )

                    divided_labels.append(
                        np.array(selected_label, dtype=np.uint8)
                    )

                    if np.shape(divided_images[-1]) != (image_size[0], image_size[1], 3) or \
                        np.shape(divided_labels[-1]) != (image_size[0], image_size[1]):
                        print(np.shape(divided_images[-1]), np.shape(divided_labels[-1]))
                        raise ValueError('Shape of image or label are inappropriate (%s, %s)' % (img, lab))

            if len(divided_images) != len(divided_labels):
                raise ValueError('Number of divided images and labels are different! (%s, %s)' % (img, lab))

            if len(divided_images) > 0:
                for d_i, d_l in zip(divided_images, divided_labels):
                    example = tf.train.Example(
                        features=tf.train.Features(
                            feature={
                                'height': int64_feature(_height),
                                'width': int64_feature(_width),
                                'image/raw': bytes_feature(d_i.tostring()),
                                'label/raw': bytes_feature(d_l.tostring())
                            }
                        )
                    )
                    if writer is not None:
                        mutex.acquire()
                        writer.write(example.SerializeToString())  # Write to tfrecord
                        mutex.release()

            mutex2.acquire()
            if timer is not None:
                timer.tick_timer()
            mutex2.release()

    timer = Timer()

    def print_function():
        sys.stdout.write('\rLoad images : [%d/%d][%.2f%%][%d hour %d minute %d second left]' %
                         (timer.current_progress, timer.whole_number_of_data,
                          timer.current_progress / timer.whole_number_of_data * 100,
                          int(timer.remain_time / 3600), int(timer.remain_time / 60) % 60, timer.remain_time % 60))
        sys.stdout.flush()

    timer.print_fn = print_function

    timer.start(len(train_list))
    run_thread_blocks(
        number_of_threads=16,
        data_list=train_list,
        thread_fn=load_image_pairs,
        writer=train_writer,
        timer=timer
    )

    timer.start(len(eval_list))
    run_thread_blocks(
        number_of_threads=16,
        data_list=eval_list,
        thread_fn=load_image_pairs,
        writer=eval_writer,
        timer=timer
    )

    print('\nCreate tfrecord [%s] finished!' % train_data_output_path)
    print('Create tfrecord [%s] finished!' % eval_data_output_path)
    train_writer.close()
    eval_writer.close()


def make_lane_tfrecord(data_path, data_output_path, shuffle_data=True):
    """
    Multi thread version of tfrecord maker
    :param data_path: Path of data
    :param data_output_path: Path of tfrecord
    :param shuffle_data: To shuffle data or not
    :return:
    """

    data_folders = file_list(data_path)

    image_list = list()
    label_list = list()
    for f in data_folders:
        _folder = file_list(f)
        image_list += file_list(_folder[1])
        label_list += file_list(_folder[0])

    image_and_label_list = [[i, j] for [i, j] in zip(image_list, label_list)]
    if shuffle_data:
        shuffle(image_and_label_list)

    print('Make tfrecord from %s to %s' % (data_path, data_output_path))
    data_writer = tf.python_io.TFRecordWriter(data_output_path)

    # Make list of image data
    print('Start make image and label pairs')

    mutex = Lock()
    mutex2 = Lock()

    image_size = (720, 1080)
    division = 1

    def load_image_pairs(thread_id, img_and_label_list, writer=None, timer=None):
        global finished_samples
        global time_mean
        global time_start
        if thread_id == 0:
            time_start = time()
            finished_samples = 0
            time_mean = -1
        _height = int(image_size[0] / division)
        _width = int(image_size[1] / division)
        for idx, (img, lab) in enumerate(img_and_label_list):
            image = np.array(Image.open(img).convert('RGB'))
            label = np.array(Image.open(lab))
            label = np.equal(label, 2).astype(np.uint8) + np.equal(label, 255).astype(np.uint8) * 255

            divided_images = list()
            divided_labels = list()
            for d_i in range(division * 2 - 1):
                for d_j in range(division * 2 - 1):
                    divided_images.append(
                        np.array(
                            image[
                                d_i * int(0.5 * _height):d_i * int(0.5 * _height) + _height,
                                d_j * int(0.5 * _width):d_j * int(0.5 * _width) + _width,
                                :3
                            ],
                            dtype=np.uint8
                        )
                    )

                    divided_labels.append(
                        np.array(
                            label[
                                d_i * int(0.5 * _height):d_i * int(0.5 * _height) + _height,
                                d_j * int(0.5 * _width):d_j * int(0.5 * _width) + _width
                            ],
                            dtype=np.uint8
                        )
                    )

                    if np.shape(divided_images[-1]) != (image_size[0], image_size[1], 3) or \
                        np.shape(divided_labels[-1]) != (image_size[0], image_size[1]):
                        print(np.shape(divided_images[-1]), np.shape(divided_labels[-1]))
                        raise ValueError('Shape of image or label are inappropriate (%s, %s)' % (img, lab))

            if len(divided_images) != len(divided_labels):
                raise ValueError('Number of divided images and labels are different! (%s, %s)' % (img, lab))

            if len(divided_images) > 0:
                for d_i, d_l in zip(divided_images, divided_labels):
                    example = tf.train.Example(
                        features=tf.train.Features(
                            feature={
                                'height': int64_feature(_height),
                                'width': int64_feature(_width),
                                'image/raw': bytes_feature(d_i.tostring()),
                                'label/raw': bytes_feature(d_l.tostring())
                            }
                        )
                    )
                    if writer is not None:
                        mutex.acquire()
                        writer.write(example.SerializeToString())  # Write to tfrecord
                        mutex.release()

            mutex2.acquire()
            if timer is not None:
                timer.tick_timer()
            mutex2.release()

    timer = Timer()

    def print_function():
        sys.stdout.write('\rLoad images : [%d/%d][%.2f%%][%d hour %d minute %d second left]' %
                         (timer.current_progress, timer.whole_number_of_data,
                          timer.current_progress / timer.whole_number_of_data * 100,
                          int(timer.remain_time / 3600), int(timer.remain_time / 60) % 60, timer.remain_time % 60))
        sys.stdout.flush()

    timer.print_fn = print_function

    timer.start(len(image_and_label_list))

    run_thread_blocks(
        number_of_threads=16,
        data_list=image_and_label_list,
        thread_fn=load_image_pairs,
        writer=data_writer,
        timer=timer
    )

    print('\nCreate tfrecord [%s] finished!' % data_output_path)
    data_writer.close()


def make_culane_tfrecord(data_path, data_output_path, shuffle_data=True):
    """
    Multi thread version of tfrecord maker
    :param data_path: Path of data
    :param data_output_path: Path of tfrecord
    :param shuffle_data: To shuffle data or not
    :return:
    """

    data_folders = file_list(data_path)
    image_folders = file_list(data_folders[0])
    label_folders = file_list(data_folders[1])

    image_list = list()
    label_list = list()
    for i_f, l_f in zip(image_folders, label_folders):
        image_list += file_list(i_f)
        label_list += file_list(l_f)

    image_and_label_list = [[i, j] for [i, j] in zip(image_list, label_list)]
    image_and_label_list = random_select(image_and_label_list, size=30000)
    if shuffle_data:
        shuffle(image_and_label_list)

    print('Make tfrecord from %s to %s' % (data_path, data_output_path))
    data_writer = tf.python_io.TFRecordWriter(data_output_path)

    # Make list of image data
    print('Start make image and label pairs')

    mutex = Lock()
    mutex2 = Lock()

    image_size = (590, 1640)
    division = 1

    def load_image_pairs(thread_id, img_and_label_list, writer=None, timer=None):
        global finished_samples
        global time_mean
        global time_start
        if thread_id == 0:
            time_start = time()
            finished_samples = 0
            time_mean = -1
        _height = int(image_size[0] / division)
        _width = int(image_size[1] / division)
        for idx, (img, lab) in enumerate(img_and_label_list):
            image = np.array(Image.open(img).convert('RGB'))
            label = np.array(Image.open(lab))
            label = np.not_equal(label, 0).astype(np.uint8)

            divided_images = list()
            divided_labels = list()
            for d_i in range(division * 2 - 1):
                for d_j in range(division * 2 - 1):
                    divided_images.append(
                        np.array(
                            image[
                                d_i * int(0.5 * _height):d_i * int(0.5 * _height) + _height,
                                d_j * int(0.5 * _width):d_j * int(0.5 * _width) + _width,
                                :3
                            ],
                            dtype=np.uint8
                        )
                    )

                    divided_labels.append(
                        np.array(
                            label[
                                d_i * int(0.5 * _height):d_i * int(0.5 * _height) + _height,
                                d_j * int(0.5 * _width):d_j * int(0.5 * _width) + _width
                            ],
                            dtype=np.uint8
                        )
                    )

                    if np.shape(divided_images[-1]) != (image_size[0], image_size[1], 3) or \
                        np.shape(divided_labels[-1]) != (image_size[0], image_size[1]):
                        print(np.shape(divided_images[-1]), np.shape(divided_labels[-1]))
                        raise ValueError('Shape of image or label are inappropriate (%s, %s)' % (img, lab))

            if len(divided_images) != len(divided_labels):
                raise ValueError('Number of divided images and labels are different! (%s, %s)' % (img, lab))

            if len(divided_images) > 0:
                for d_i, d_l in zip(divided_images, divided_labels):
                    example = tf.train.Example(
                        features=tf.train.Features(
                            feature={
                                'height': int64_feature(_height),
                                'width': int64_feature(_width),
                                'image/raw': bytes_feature(d_i.tostring()),
                                'label/raw': bytes_feature(d_l.tostring())
                            }
                        )
                    )
                    if writer is not None:
                        mutex.acquire()
                        writer.write(example.SerializeToString())  # Write to tfrecord
                        mutex.release()

            mutex2.acquire()
            if timer is not None:
                timer.tick_timer()
            mutex2.release()

    timer = Timer()

    def print_function():
        sys.stdout.write('\rLoad images : [%d/%d][%.2f%%][%d hour %d minute %d second left]' %
                         (timer.current_progress, timer.whole_number_of_data,
                          timer.current_progress / timer.whole_number_of_data * 100,
                          int(timer.remain_time / 3600), int(timer.remain_time / 60) % 60, timer.remain_time % 60))
        sys.stdout.flush()

    timer.print_fn = print_function

    timer.start(len(image_and_label_list))

    run_thread_blocks(
        number_of_threads=16,
        data_list=image_and_label_list,
        thread_fn=load_image_pairs,
        writer=data_writer,
        timer=timer
    )

    print('\nCreate tfrecord [%s] finished!' % data_output_path)
    data_writer.close()


def make_bdd100k_tfrecord(data_path, label_path, data_output_path, shuffle_data=True):
    """
    Multi thread version of tfrecord maker
    :param data_path: Path of data
    :param data_output_path: Path of tfrecord
    :param shuffle_data: To shuffle data or not
    :return:
    """

    print('[Lead json]')
    data = open(label_path).read()
    print('[Done]')

    print('[Load json]')
    data = json.loads(data)
    print('[Done]')

    print('[Gather data]')
    dataset = list()
    for d in data:
        vertices = list()
        closed = list()
        for l in d['labels']:
            if l['category'] == 'lane' and (
                    l['attributes']['laneType'] == 'single white' or
                    l['attributes']['laneType'] == 'single yellow' or
                    l['attributes']['laneType'] == 'double white' or
                    l['attributes']['laneType'] == 'double yellow'
            ):
                if len(l['poly2d']) != 1:
                    raise ValueError('Poly2d!!', d['name'])
                vertices.append(l['poly2d'][0]['vertices'])
                closed.append(l['poly2d'][0]['closed'])

        if len(vertices) == 0:
            continue

        _d = {
            'image_path': os.path.join(data_path, d['name']),
            'vertices': vertices,
            'closed': closed
        }
        dataset.append(_d)
    print('[Done]')

    print('[Random select and shuffle data]')
    dataset = random_select(dataset, 30000)
    if shuffle_data:
        shuffle(dataset)
    print('[Done]')
    print('Number of data', len(dataset))

    print('Make tfrecord from %s to %s' % (data_path, data_output_path))
    data_writer = tf.python_io.TFRecordWriter(data_output_path)

    # Make list of image data
    print('Start make image and label pairs')

    mutex = Lock()
    mutex2 = Lock()

    image_size = (720, 1280)
    division = 1

    def load_image_pairs(thread_id, img_and_label_list, writer=None, timer=None):
        global finished_samples
        global time_mean
        global time_start
        if thread_id == 0:
            time_start = time()
            finished_samples = 0
            time_mean = -1
        _height = int(image_size[0] / division)
        _width = int(image_size[1] / division)
        for idx, _data in enumerate(img_and_label_list):
            image = np.array(Image.open(_data['image_path']).convert('RGB'))
            label = np.zeros_like(image, dtype=np.uint8)
            for v, c in zip(_data['vertices'], _data['closed']):
                cv2.polylines(label, np.array([v], dtype=np.int32), c, 1, thickness=12)

            divided_images = list()
            divided_labels = list()
            for d_i in range(division * 2 - 1):
                for d_j in range(division * 2 - 1):
                    divided_images.append(
                        np.array(
                            image[
                                d_i * int(0.5 * _height):d_i * int(0.5 * _height) + _height,
                                d_j * int(0.5 * _width):d_j * int(0.5 * _width) + _width,
                                :3
                            ],
                            dtype=np.uint8
                        )
                    )

                    divided_labels.append(
                        np.array(
                            label[
                                d_i * int(0.5 * _height):d_i * int(0.5 * _height) + _height,
                                d_j * int(0.5 * _width):d_j * int(0.5 * _width) + _width
                            ],
                            dtype=np.uint8
                        )
                    )

                    if np.shape(divided_images[-1]) != (image_size[0], image_size[1], 3) or \
                        np.shape(divided_labels[-1]) != (image_size[0], image_size[1]):
                        print(np.shape(divided_images[-1]), np.shape(divided_labels[-1]))
                        raise ValueError('Shape of image or label are inappropriate (%s)' % (_data['image_path']))

            if len(divided_images) != len(divided_labels):
                raise ValueError('Number of divided images and labels are different! (%s)' % (_data['image_path']))

            if len(divided_images) > 0:
                for d_i, d_l in zip(divided_images, divided_labels):
                    example = tf.train.Example(
                        features=tf.train.Features(
                            feature={
                                'height': int64_feature(_height),
                                'width': int64_feature(_width),
                                'image/raw': bytes_feature(d_i.tostring()),
                                'label/raw': bytes_feature(d_l.tostring())
                            }
                        )
                    )
                    if writer is not None:
                        mutex.acquire()
                        writer.write(example.SerializeToString())  # Write to tfrecord
                        mutex.release()

            mutex2.acquire()
            if timer is not None:
                timer.tick_timer()
            mutex2.release()

    timer = Timer()

    def print_function():
        sys.stdout.write('\rLoad images : [%d/%d][%.2f%%][%d hour %d minute %d second left]' %
                         (timer.current_progress, timer.whole_number_of_data,
                          timer.current_progress / timer.whole_number_of_data * 100,
                          int(timer.remain_time / 3600), int(timer.remain_time / 60) % 60, timer.remain_time % 60))
        sys.stdout.flush()

    timer.print_fn = print_function

    timer.start(len(dataset))

    run_thread_blocks(
        number_of_threads=16,
        data_list=dataset,
        thread_fn=load_image_pairs,
        writer=data_writer,
        timer=timer
    )

    print('\nCreate tfrecord [%s] finished!' % data_output_path)
    data_writer.close()


def main(unused_argv):

    make_bdd100k_tfrecord(
        data_path='D:/BDD100k/bdd100k/images/100k/train',
        label_path='D:/BDD100k/bdd100k_labels_images_train.json',
        data_output_path='D:/BDD100k/lane_train.tfrecord',
        shuffle_data=True
    )

    # make_culane_tfrecord(
    #     data_path='D:/CULane',
    #     data_output_path='D:/CULane/lane_train.tfrecord',
    #     shuffle_data=True
    # )

    # make_lane_tfrecord(
    #     data_path='D:/HighwayDriving/Train',
    #     data_output_path='C:/Users/YDK/Desktop/YDK/Graduate School/Work/STAR Lab/LaneSegmentationNetwork' +
    #                      '/dataset_lane/lane_train.tfrecord',
    #     shuffle_data=True
    # )
    # make_lane_tfrecord(
    #     data_path='D:/HighwayDriving/Test',
    #     data_output_path='C:/Users/YDK/Desktop/YDK/Graduate School/Work/STAR Lab/LaneSegmentationNetwork' +
    #                      '/dataset_lane/lane_test.tfrecord',
    #     shuffle_data=True
    # )
    # make_shi_tfrecord(image_path='D:/Samsung Heavy Industries/DB/original',
    #                   label_path='D:/Samsung Heavy Industries/DB/LabelsConverted',
    #                   train_data_output_path='D:/Samsung Heavy Industries/DB/LabelsConverted/' +
    #                                          'SHI_dataset_ori.tfrecord',
    #                   eval_data_output_path='D:/Samsung Heavy Industries/DB/LabelsConverted/' +
    #                                         'SHI_dataset_eval_ori.tfrecord',
    #                   shuffle_data=True)


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