"""
This file is used to generate TFRecords using the AI Challenger dataset.

@ Author: <NAME>. <EMAIL>

@ Date created: Jun 04, 2019

@ Last modified: Jun 27, 2019

"""
import numpy as np
import os
import time
from skimage import io, draw
from skimage.transform import resize
from random import shuffle
import tensorflow as tf
import json

"""
# The image which contains a person is collected from the AI Challenger dataset in the following steps:
    1. Get the coordinate of the bounding box in the original image.
    2. Adjust the ratio of the bounding box to be 4:3 (height : width)

    Note that the coordinates of keypoints are also re-calculated when the foreground parts are clipped from the 
    original images.

"""


def draw_points_on_img(img, point_ver, point_hor, point_class):
    for i in range(len(point_class)):
        if point_class[i] != 3:
            rr, cc = draw.circle(point_ver[i], point_hor[i], 10, (256, 192))
            #draw.set_color(img, [rr, cc], [0., 0., 0.], alpha=5)
            img[rr, cc, :] = 0
    #io.imshow(img)
    #io.show()

    return img


def draw_lines_on_img(img, point_ver, point_hor, point_class):
    line_list = [[0, 1], [1, 2], [3, 4], [4, 5], [6, 7], [7, 8], [9, 10],
               [10, 11], [12, 13], [13, 6], [13, 9], [13, 0], [13, 3]]

    # key point class: 1:visible, 2: not visible, 3: not marked
    for start_point_id in range(len(point_class)):
        if point_class[start_point_id] == 3:
            continue
        for end_point_id in range(len(point_class)):
            if point_class[end_point_id] == 3:
                continue

            if [start_point_id, end_point_id] in line_list:
                rr, cc = draw.line(int(point_ver[start_point_id]), int(point_hor[start_point_id]),
                                   int(point_ver[end_point_id]), int(point_hor[end_point_id]))
                draw.set_color(img, [rr, cc], [255, 0, 0])

    return img


def _int64_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 _float_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))


def extract_people_from_dataset(dataset_root_path='../../../dataset/ai_challenger/', image_save_path='../dataset/imgs/',
                                tfrecords_path='../dataset/', is_shuffle=True):
    """
    This function is used to extract people from the AI Challenger dataset. The extract image will contain only one
    person each and will be saved as a single .jpg file. At last, the image and the the corresponding annotation will
    be saved into a .tfrecord file.

    :param dataset_root_path: the root path of the AI Challenger dataset.
    :param image_save_path: the path used to save the clipped images.
    :param tfrecord_path: the path used to save the .tfrecords file.
    :param is_shuffle: is shuffle.
    :return: None.
    """
    annotation_file = os.path.join(dataset_root_path, 'keypoint_train_annotations_20170909.json')
    image_read_path = os.path.join(dataset_root_path, 'train_images')
    tfrecords_file = os.path.join(tfrecords_path, 'train.tfrecords')

    if not os.path.exists(tfrecords_path):
        os.mkdir(tfrecords_path)
    if os.path.exists(tfrecords_file):
        os.remove(tfrecords_file)
    if os.path.exists(image_save_path):
        useless = os.listdir(image_save_path)
        for onefile in useless:
            os.remove(os.path.join(image_save_path, onefile))
    else:
        os.mkdir(image_save_path)

    saved_number = 0
    image_number = 0
    start_time = time.time()
    with tf.python_io.TFRecordWriter(tfrecords_file) as tfwriter:

        with open(annotation_file, 'r') as jsfile:
            data = json.load(jsfile)

            for one_item in data:
                img_id = one_item['image_id']
                image_number += 1
                if image_number % 100 == 0:
                    print('Processed {} images, extracted {} people from the dataset. '
                          'time = {}'.format(image_number, saved_number, time.time() - start_time))

                kps = one_item['keypoint_annotations']
                boxes = one_item['human_annotations']

                # read image
                img_filename = os.path.join(image_read_path, img_id + '.jpg')
                img = io.imread(img_filename)

                for i in range(len(boxes)):
                    # construct the name of a human in the dictionary,
                    # for example, the first one (when i = 0) is 'human1'
                    human_name = 'human' + str(i+1)

                    kp = kps[human_name]
                    box = boxes[human_name]
                    p1_hor, p1_ver, p2_hor, p2_ver = box
                    foreground = img[p1_ver:p2_ver, p1_hor:p2_hor, :]

                    try:
                        foreground = resize(foreground, (256, 192, 3))
                    except ValueError:
                        print('ValueError at image {} and {}'.format(image_number, human_name))
                        continue

                    foreground = foreground * 255.0
                    foreground_uint8 = np.uint8(foreground)

                    kp_hor = (np.array(kp[0::3]) - p1_hor) / (p2_hor - p1_hor) * 192
                    kp_ver = (np.array(kp[1::3]) - p1_ver) / (p2_ver - p1_ver) * 256
                    kp_class = np.array(kp[2::3])

                    img_name = img_id + '_' + human_name + '.jpg'

                    io.imsave(os.path.join(image_save_path, img_id + '_' + human_name + '.jpg'), foreground_uint8)

                    example = tf.train.Example(
                        features=tf.train.Features(
                            feature={
                                'image_name': _bytes_feature(img_name.encode()),
                                'image_raw': _bytes_feature(foreground_uint8.tobytes()),
                                'keypoints_ver': _bytes_feature(np.uint8(kp_ver).tobytes()),
                                'keypoints_hor': _bytes_feature(np.uint8(kp_hor).tobytes()),
                                'keypoints_class': _bytes_feature(np.uint8(kp_class).tobytes())
                            }))
                    tfwriter.write(example.SerializeToString())

                    saved_number += 1
    print('Extracted {} people from the dataset in total.'.format(saved_number))


def decode_proto(proto):
    features = tf.parse_single_example(proto,
                                       features={
                                           'image_name': tf.FixedLenFeature([], tf.string),
                                           'image_raw': tf.FixedLenFeature([], tf.string),
                                           'keypoints_ver': tf.FixedLenFeature([], tf.string),
                                           'keypoints_hor': tf.FixedLenFeature([], tf.string),
                                           'keypoints_class': tf.FixedLenFeature([], tf.string),
                                       })
    image_name = features['image_name']

    image_raw = tf.decode_raw(features['image_raw'], out_type=np.uint8)
    image = tf.reshape(image_raw, [256, 192, 3])

    keypoints_ver = tf.decode_raw(features['keypoints_ver'], out_type=np.uint8)
    keypoints_hor = tf.decode_raw(features['keypoints_hor'], out_type=np.uint8)
    keypoints_class = tf.decode_raw(features['keypoints_class'], out_type=np.uint8)
    return image_name, image, keypoints_ver, keypoints_hor, keypoints_class


def decode_tfrecord(filename_queue):
    tfreader = tf.TFRecordReader()
    _, proto = tfreader.read(filename_queue)
    image_name, image, keypoints_ver, keypoints_hor, keypoints_class = decode_proto(proto)

    return image_name, image, keypoints_ver, keypoints_hor, keypoints_class


def input_batch(datasetname, batch_size, num_epochs):
    """
    This function is used to decode the TFrecord and return a batch of images as well as their information
    :param datasetname: the name of the TFrecord file.
    :param batch_size: the number of images in a batch
    :param num_epochs: the number of epochs
    :return: a batch of images as well as their information
    """
    with tf.name_scope('input_batch'):
        # The shuffle transformation uses a finite-sized buffer to shuffle elements
        # in memory. The parameter is the number of elements in the buffer. For
        # completely uniform shuffling, set the parameter to be the same as the
        # number of elements in the dataset.
        mydataset = tf.data.TFRecordDataset(datasetname)
        mydataset = mydataset.map(decode_proto)

        # have no idea why I can't set the parameter of mydataset.shuffle to be the number of the dataset......
        # mydataset = mydataset.shuffle(200)
        mydataset = mydataset.repeat(num_epochs * 2)
        # drop all the data that can't be used to make up a batch
        mydataset = mydataset.batch(batch_size, drop_remainder=True)
        iterator = mydataset.make_one_shot_iterator()

        nextelement = iterator.get_next()
        return nextelement


def mytest():
    tfrecord_file = '../dataset/train.tfrecords'

    filename_queue = tf.train.string_input_producer([tfrecord_file], num_epochs=None)
    image_name, image, keypoints_ver, keypoints_hor, keypoints_class = decode_tfrecord(filename_queue)

    with tf.Session() as sess:
        init_op = tf.global_variables_initializer()
        sess.run(init_op)
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        try:
            # while not coord.should_stop():
            for i in range(10):
                img_name, img, point_ver, point_hor, point_class = sess.run([image_name, image, keypoints_ver,
                                                                             keypoints_hor, keypoints_class])

                print(img_name, point_hor, point_ver, point_class)

                for i in range(len(point_class)):
                    if point_class[i] > 0:
                        rr, cc = draw.circle(point_ver[i], point_hor[i], 10, (256, 192))
                        img[rr, cc, :] = 0

                io.imshow(img)
                io.show()

        except tf.errors.OutOfRangeError:
            print('Done reading')
        finally:
            coord.request_stop()


if __name__ == '__main__':
    extract_people_from_dataset()
    #mytest()



"""
This file is used to generate TFRecords using the AI Challenger dataset.

@ Author: <NAME>. <EMAIL>

@ Date created: Jun 04, 2019

@ Last modified: Jun 27, 2019

"""
import numpy as np
import os
import time
from skimage import io, draw
from skimage.transform import resize
from random import shuffle
import tensorflow as tf
import json

"""
# The image which contains a person is collected from the AI Challenger dataset in the following steps:
    1. Get the coordinate of the bounding box in the original image.
    2. Adjust the ratio of the bounding box to be 4:3 (height : width)

    Note that the coordinates of keypoints are also re-calculated when the foreground parts are clipped from the 
    original images.

"""


def draw_points_on_img(img, point_ver, point_hor, point_class):
    for i in range(len(point_class)):
        if point_class[i] != 3:
            rr, cc = draw.circle(point_ver[i], point_hor[i], 10, (256, 192))
            #draw.set_color(img, [rr, cc], [0., 0., 0.], alpha=5)
            img[rr, cc, :] = 0
    #io.imshow(img)
    #io.show()

    return img


def draw_lines_on_img(img, point_ver, point_hor, point_class):
    line_list = [[0, 1], [1, 2], [3, 4], [4, 5], [6, 7], [7, 8], [9, 10],
               [10, 11], [12, 13], [13, 6], [13, 9], [13, 0], [13, 3]]

    # key point class: 1:visible, 2: not visible, 3: not marked
    for start_point_id in range(len(point_class)):
        if point_class[start_point_id] == 3:
            continue
        for end_point_id in range(len(point_class)):
            if point_class[end_point_id] == 3:
                continue

            if [start_point_id, end_point_id] in line_list:
                rr, cc = draw.line(int(point_ver[start_point_id]), int(point_hor[start_point_id]),
                                   int(point_ver[end_point_id]), int(point_hor[end_point_id]))
                draw.set_color(img, [rr, cc], [255, 0, 0])

    return img


def _int64_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 _float_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))


def extract_people_from_dataset(dataset_root_path='../../../dataset/ai_challenger/', image_save_path='../dataset/imgs/',
                                tfrecords_path='../dataset/', is_shuffle=True):
    """
    This function is used to extract people from the AI Challenger dataset. The extract image will contain only one
    person each and will be saved as a single .jpg file. At last, the image and the the corresponding annotation will
    be saved into a .tfrecord file.

    :param dataset_root_path: the root path of the AI Challenger dataset.
    :param image_save_path: the path used to save the clipped images.
    :param tfrecord_path: the path used to save the .tfrecords file.
    :param is_shuffle: is shuffle.
    :return: None.
    """
    annotation_file = os.path.join(dataset_root_path, 'keypoint_train_annotations_20170909.json')
    image_read_path = os.path.join(dataset_root_path, 'train_images')
    tfrecords_file = os.path.join(tfrecords_path, 'train.tfrecords')

    if not os.path.exists(tfrecords_path):
        os.mkdir(tfrecords_path)
    if os.path.exists(tfrecords_file):
        os.remove(tfrecords_file)
    if os.path.exists(image_save_path):
        useless = os.listdir(image_save_path)
        for onefile in useless:
            os.remove(os.path.join(image_save_path, onefile))
    else:
        os.mkdir(image_save_path)

    saved_number = 0
    image_number = 0
    start_time = time.time()
    with tf.python_io.TFRecordWriter(tfrecords_file) as tfwriter:

        with open(annotation_file, 'r') as jsfile:
            data = json.load(jsfile)

            for one_item in data:
                img_id = one_item['image_id']
                image_number += 1
                if image_number % 100 == 0:
                    print('Processed {} images, extracted {} people from the dataset. '
                          'time = {}'.format(image_number, saved_number, time.time() - start_time))

                kps = one_item['keypoint_annotations']
                boxes = one_item['human_annotations']

                # read image
                img_filename = os.path.join(image_read_path, img_id + '.jpg')
                img = io.imread(img_filename)

                for i in range(len(boxes)):
                    # construct the name of a human in the dictionary,
                    # for example, the first one (when i = 0) is 'human1'
                    human_name = 'human' + str(i+1)

                    kp = kps[human_name]
                    box = boxes[human_name]
                    p1_hor, p1_ver, p2_hor, p2_ver = box
                    foreground = img[p1_ver:p2_ver, p1_hor:p2_hor, :]

                    try:
                        foreground = resize(foreground, (256, 192, 3))
                    except ValueError:
                        print('ValueError at image {} and {}'.format(image_number, human_name))
                        continue

                    foreground = foreground * 255.0
                    foreground_uint8 = np.uint8(foreground)

                    kp_hor = (np.array(kp[0::3]) - p1_hor) / (p2_hor - p1_hor) * 192
                    kp_ver = (np.array(kp[1::3]) - p1_ver) / (p2_ver - p1_ver) * 256
                    kp_class = np.array(kp[2::3])

                    img_name = img_id + '_' + human_name + '.jpg'

                    io.imsave(os.path.join(image_save_path, img_id + '_' + human_name + '.jpg'), foreground_uint8)

                    example = tf.train.Example(
                        features=tf.train.Features(
                            feature={
                                'image_name': _bytes_feature(img_name.encode()),
                                'image_raw': _bytes_feature(foreground_uint8.tobytes()),
                                'keypoints_ver': _bytes_feature(np.uint8(kp_ver).tobytes()),
                                'keypoints_hor': _bytes_feature(np.uint8(kp_hor).tobytes()),
                                'keypoints_class': _bytes_feature(np.uint8(kp_class).tobytes())
                            }))
                    tfwriter.write(example.SerializeToString())

                    saved_number += 1
    print('Extracted {} people from the dataset in total.'.format(saved_number))


def decode_proto(proto):
    features = tf.parse_single_example(proto,
                                       features={
                                           'image_name': tf.FixedLenFeature([], tf.string),
                                           'image_raw': tf.FixedLenFeature([], tf.string),
                                           'keypoints_ver': tf.FixedLenFeature([], tf.string),
                                           'keypoints_hor': tf.FixedLenFeature([], tf.string),
                                           'keypoints_class': tf.FixedLenFeature([], tf.string),
                                       })
    image_name = features['image_name']

    image_raw = tf.decode_raw(features['image_raw'], out_type=np.uint8)
    image = tf.reshape(image_raw, [256, 192, 3])

    keypoints_ver = tf.decode_raw(features['keypoints_ver'], out_type=np.uint8)
    keypoints_hor = tf.decode_raw(features['keypoints_hor'], out_type=np.uint8)
    keypoints_class = tf.decode_raw(features['keypoints_class'], out_type=np.uint8)
    return image_name, image, keypoints_ver, keypoints_hor, keypoints_class


def decode_tfrecord(filename_queue):
    tfreader = tf.TFRecordReader()
    _, proto = tfreader.read(filename_queue)
    image_name, image, keypoints_ver, keypoints_hor, keypoints_class = decode_proto(proto)

    return image_name, image, keypoints_ver, keypoints_hor, keypoints_class


def input_batch(datasetname, batch_size, num_epochs):
    """
    This function is used to decode the TFrecord and return a batch of images as well as their information
    :param datasetname: the name of the TFrecord file.
    :param batch_size: the number of images in a batch
    :param num_epochs: the number of epochs
    :return: a batch of images as well as their information
    """
    with tf.name_scope('input_batch'):
        # The shuffle transformation uses a finite-sized buffer to shuffle elements
        # in memory. The parameter is the number of elements in the buffer. For
        # completely uniform shuffling, set the parameter to be the same as the
        # number of elements in the dataset.
        mydataset = tf.data.TFRecordDataset(datasetname)
        mydataset = mydataset.map(decode_proto)

        # have no idea why I can't set the parameter of mydataset.shuffle to be the number of the dataset......
        # mydataset = mydataset.shuffle(200)
        mydataset = mydataset.repeat(num_epochs * 2)
        # drop all the data that can't be used to make up a batch
        mydataset = mydataset.batch(batch_size, drop_remainder=True)
        iterator = mydataset.make_one_shot_iterator()

        nextelement = iterator.get_next()
        return nextelement


def mytest():
    tfrecord_file = '../dataset/train.tfrecords'

    filename_queue = tf.train.string_input_producer([tfrecord_file], num_epochs=None)
    image_name, image, keypoints_ver, keypoints_hor, keypoints_class = decode_tfrecord(filename_queue)

    with tf.Session() as sess:
        init_op = tf.global_variables_initializer()
        sess.run(init_op)
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        try:
            # while not coord.should_stop():
            for i in range(10):
                img_name, img, point_ver, point_hor, point_class = sess.run([image_name, image, keypoints_ver,
                                                                             keypoints_hor, keypoints_class])

                print(img_name, point_hor, point_ver, point_class)

                for i in range(len(point_class)):
                    if point_class[i] > 0:
                        rr, cc = draw.circle(point_ver[i], point_hor[i], 10, (256, 192))
                        img[rr, cc, :] = 0

                io.imshow(img)
                io.show()

        except tf.errors.OutOfRangeError:
            print('Done reading')
        finally:
            coord.request_stop()


if __name__ == '__main__':
    extract_people_from_dataset()
    #mytest()



This file is used to generate TFRecords using the AI Challenger dataset.  @ Author: <NAME>. <EMAIL>  @ Date created: Jun 04, 2019  @ Last modified: Jun 27, 2019 # The image which contains a person is collected from the AI Challenger dataset in the following steps:     1. Get the coordinate of the bounding box in the original image.     2. Adjust the ratio of the bounding box to be 4:3 (height : width)      Note that the coordinates of keypoints are also re-calculated when the foreground parts are clipped from the      original images. #draw.set_color(img, [rr, cc], [0., 0., 0.], alpha=5) #io.imshow(img) #io.show() # key point class: 1:visible, 2: not visible, 3: not marked This function is used to extract people from the AI Challenger dataset. The extract image will contain only one     person each and will be saved as a single .jpg file. At last, the image and the the corresponding annotation will     be saved into a .tfrecord file.      :param dataset_root_path: the root path of the AI Challenger dataset.     :param image_save_path: the path used to save the clipped images.     :param tfrecord_path: the path used to save the .tfrecords file.     :param is_shuffle: is shuffle.     :return: None. # read image # construct the name of a human in the dictionary, # for example, the first one (when i = 0) is 'human1' This function is used to decode the TFrecord and return a batch of images as well as their information     :param datasetname: the name of the TFrecord file.     :param batch_size: the number of images in a batch     :param num_epochs: the number of epochs     :return: a batch of images as well as their information # The shuffle transformation uses a finite-sized buffer to shuffle elements # in memory. The parameter is the number of elements in the buffer. For # completely uniform shuffling, set the parameter to be the same as the # number of elements in the dataset. # have no idea why I can't set the parameter of mydataset.shuffle to be the number of the dataset...... # mydataset = mydataset.shuffle(200) # drop all the data that can't be used to make up a batch # while not coord.should_stop(): #mytest()