# Path: utils/camvid.py
# IMAGE_HEIGHT = 720
# IMAGE_WIDTH = 960
# IMAGE_DEPTH = 3
# NUM_CLASSES = 32
# NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN = 580
# NUM_EXAMPLES_PER_EPOCH_FOR_TEST = 580
# NUM_EXAMPLES_PER_EPOCH_FOR_EVAL = 580
# def _generate_image_and_label_batch(image, label, min_queue_examples,
#                                     batch_size, shuffle):
# def CamVid_reader_seq(filename_queue, seq_length):
# def CamVid_reader(filename_queue):
# def get_filename_list(path):
# def CamVidInputs(image_filenames, label_filenames, batch_size, shuffle=True):


import os
import tensorflow as tf
import segnet_vgg
from utils import camvid
NUM_CLASSES = 32
INITIAL_LEARNING_RATE = 0.0001

image_dir = "/usr/local/google/home/limeng/Downloads/camvid/train.txt"
val_dir = "/usr/local/google/home/limeng/Downloads/camvid/val.txt"


def loss(logits, labels):
    logits = tf.reshape(logits, [-1, NUM_CLASSES])
    labels = tf.reshape(labels, [-1])

    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=logits, labels=labels, name='cross_entropy')
    cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy_mean')
    tf.add_to_collection('losses', cross_entropy_mean)
    return tf.add_n(tf.get_collection('losses'), name='total_loss')


def train(total_loss):
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        global_step = tf.Variable(0, name='global_step', trainable=False)
        learning_rate = tf.train.exponential_decay(
            INITIAL_LEARNING_RATE, global_step, EPOCH * 0.2, 0.9, staircase=True)
        tf.summary.scalar('total_loss', total_loss)
        optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
        return optimizer.minimize(total_loss, global_step=global_step)


def main(_):

# Path: utils/camvid.py
# IMAGE_HEIGHT = 720
# IMAGE_WIDTH = 960
# IMAGE_DEPTH = 3
# NUM_CLASSES = 32
# NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN = 580
# NUM_EXAMPLES_PER_EPOCH_FOR_TEST = 580
# NUM_EXAMPLES_PER_EPOCH_FOR_EVAL = 580
# def _generate_image_and_label_batch(image, label, min_queue_examples,
#                                     batch_size, shuffle):
# def CamVid_reader_seq(filename_queue, seq_length):
# def CamVid_reader(filename_queue):
# def get_filename_list(path):
# def CamVidInputs(image_filenames, label_filenames, batch_size, shuffle=True):


# Path: self_driving/segnet/train.py
import os
import tensorflow as tf
import segnet_vgg
from utils import camvid

NUM_CLASSES = 32
INITIAL_LEARNING_RATE = 0.0001

image_dir = "/usr/local/google/home/limeng/Downloads/camvid/train.txt"
val_dir = "/usr/local/google/home/limeng/Downloads/camvid/val.txt"


def loss(logits, labels):
    logits = tf.reshape(logits, [-1, NUM_CLASSES])
    labels = tf.reshape(labels, [-1])

    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=logits, labels=labels, name='cross_entropy')
    cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy_mean')
    tf.add_to_collection('losses', cross_entropy_mean)
    return tf.add_n(tf.get_collection('losses'), name='total_loss')


def train(total_loss):
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        global_step = tf.Variable(0, name='global_step', trainable=False)
        learning_rate = tf.train.exponential_decay(
            INITIAL_LEARNING_RATE, global_step, EPOCH * 0.2, 0.9, staircase=True)
        tf.summary.scalar('total_loss', total_loss)
        optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
        return optimizer.minimize(total_loss, global_step=global_step)


def main(_):

    image_filenames, label_filenames = camvid.get_filename_list(image_dir)