# Copyright 2016 Google Inc. 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.
# ==============================================================================
#Copyright 2017 smartjinyu

#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.

import random
import threading

import numpy as np
import tensorflow as tf
from PIL import Image
import os

tf.app.flags.DEFINE_string('train_directory', 'D:\\\Python\\Small_Contribution\\image\\Train',
                           'Training data directory')
tf.app.flags.DEFINE_string('validation_directory', 'D:\\\Python\\Small_Contribution\\image\\Validate',
                           'Validation data directory')
tf.app.flags.DEFINE_string('test_directory', 'D:\\\Python\\Small_Contribution\\image\\Test',
                           'Test data directory')
tf.app.flags.DEFINE_string('output_directory', 'D:\\\Python\\Small_Contribution\\TFRecords',
                           'Output data directory')

tf.app.flags.DEFINE_integer('train_shards', 10,
                            'Number of shards in training TFRecord files.')
tf.app.flags.DEFINE_integer('validation_shards', 2,
                            'Number of shards in validation TFRecord files.')
tf.app.flags.DEFINE_integer('test_shards', 2,
                            'Number of shards in test TFRecord files.')

tf.app.flags.DEFINE_integer('num_threads', 5,
                            'Number of threads to preprocess the images.')

tf.app.flags.DEFINE_string('labels_file', 'D:\\\Python\\Small_Contribution\\image\\image.txt', 'Labels file')

FLAGS = tf.app.flags.FLAGS

def _int64_feature(value):
    if not isinstance(value, list):
        value = [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 _convert_to_example(image_buffer, label):
    example = tf.train.Example(features=tf.train.Features(feature={
        '_label': _int64_feature(label),
        '_images': _bytes_feature(tf.compat.as_bytes(image_buffer))}))
    return example

def _process_image(filename):
    image = Image.open(filename)
    image_data = image.tobytes()
    return image_data

def _process_image_files_batch(thread_index, ranges, name, filenames,
                               labels, num_shards):
    num_threads = len(ranges)
    assert not num_shards % num_threads
    num_shards_per_batch = int(num_shards / num_threads)

    shard_ranges = np.linspace(ranges[thread_index][0],
                               ranges[thread_index][1],
                               num_shards_per_batch + 1).astype(int)

    counter = 0
    for s in range(num_shards_per_batch):
        shard = thread_index * num_shards_per_batch + s
        output_filename = '%s-%.5d-of-%.5d' % (name, shard, num_shards)
        output_file = os.path.join(FLAGS.output_directory, output_filename)
        writer = tf.python_io.TFRecordWriter(output_file)

        shard_counter = 0
        files_in_shard = np.arange(shard_ranges[s], shard_ranges[s + 1], dtype=int)
        for i in files_in_shard:
            filename = filenames[i]
            label = labels[i]
            image_buffer= _process_image(filename)

            example = _convert_to_example(image_buffer, label)
            writer.write(example.SerializeToString())
            shard_counter += 1
            counter += 1

        writer.close()

def _process_image_files(name, filenames, labels, num_shards):
    spacing = np.linspace(0, len(filenames), FLAGS.num_threads + 1).astype(np.int)
    ranges = []
    for i in range(len(spacing) - 1):
        ranges.append([spacing[i], spacing[i + 1]])

    coord = tf.train.Coordinator()

    threads = []
    for thread_index in range(len(ranges)):
        args = (thread_index, ranges, name, filenames,
                labels, num_shards)
        t = threading.Thread(target=_process_image_files_batch, args=args)
        t.start()
        threads.append(t)

    coord.join(threads)

def _find_image_files(data_dir, labels_file):
    unique_labels = [l.strip() for l in tf.gfile.FastGFile(
        labels_file, 'r').readlines()]

    labels = []
    filenames = []

    label_index = 1

    for text in unique_labels:
        jpeg_file_path = '%s/%s/*' % (data_dir, text)
        matching_files = tf.gfile.Glob(jpeg_file_path)

        labels.extend([label_index] * len(matching_files))
        filenames.extend(matching_files)

        label_index += 1

    shuffled_index = list(range(len(filenames)))
    random.seed(12345)
    random.shuffle(shuffled_index)

    filenames = [filenames[i] for i in shuffled_index]
    labels = [labels[i] for i in shuffled_index]

    return filenames, labels

def _process_dataset(name, directory, num_shards, labels_file):
    filenames, labels = _find_image_files(directory, labels_file)
    _process_image_files(name, filenames, labels, num_shards)

def main(unused_argv):
    if not os.path.exists(FLAGS.output_directory):
        os.makedirs(FLAGS.output_directory)

    _process_dataset('validation', FLAGS.validation_directory,
                     FLAGS.validation_shards, FLAGS.labels_file)
    _process_dataset('train', FLAGS.train_directory,
                     FLAGS.train_shards, FLAGS.labels_file)
    _process_dataset('test', FLAGS.test_directory,
                     FLAGS.test_shards, FLAGS.labels_file)

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