#!/usr/bin/env python3
import os.path
import tensorflow as tf
import helper
import warnings
from distutils.version import LooseVersion
import project_tests as tests


# Check TensorFlow Version
assert LooseVersion(tf.__version__) >= LooseVersion('1.0'), 'Please use TensorFlow version 1.0 or newer.  You are using {}'.format(tf.__version__)
print('TensorFlow Version: {}'.format(tf.__version__))

# Check for a GPU
if not tf.test.gpu_device_name():
    warnings.warn('No GPU found. Please use a GPU to train your neural network.')
else:
    print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))


def load_vgg(sess, vgg_path):
    """
    Load Pretrained VGG Model into TensorFlow.
    :param sess: TensorFlow Session
    :param vgg_path: Path to vgg folder, containing "variables/" and "saved_model.pb"
    :return: Tuple of Tensors from VGG model (image_input, keep_prob, layer3_out, layer4_out, layer7_out)
    """
    tf.saved_model.loader.load(sess, ['vgg16'], vgg_path)
    graph = tf.get_default_graph()
    image_input = graph.get_tensor_by_name('image_input:0')
    keep_prob = graph.get_tensor_by_name('keep_prob:0')
    layer3_out = graph.get_tensor_by_name('layer3_out:0')
    layer4_out = graph.get_tensor_by_name('layer4_out:0')
    layer7_out = graph.get_tensor_by_name('layer7_out:0')

    return image_input, keep_prob, layer3_out, layer4_out, layer7_out
tests.test_load_vgg(load_vgg, tf)


def layers(vgg_layer3_out, vgg_layer4_out, vgg_layer7_out, num_classes):
    """
    Create the layers for a fully convolutional network.  Build skip-layers using the vgg layers.
    :param vgg_layer3_out: TF Tensor for VGG Layer 3 output
    :param vgg_layer4_out: TF Tensor for VGG Layer 4 output
    :param vgg_layer7_out: TF Tensor for VGG Layer 7 output
    :param num_classes: Number of classes to classify
    :return: The Tensor for the last layer of output
    """
    # 1 by 1 convolution
    layer3_1_1 = tf.layers.conv2d(vgg_layer3_out, num_classes,
        kernel_size=1, padding="same", strides=1,
        kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))
    # scaling
    layer3_1_1 = tf.multiply(layer3_1_1, 0.0001)

    layer4_1_1 = tf.layers.conv2d(vgg_layer4_out, num_classes,
        kernel_size=1, padding="same", strides=1,
        kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))
    # scaling
    layer4_1_1 = tf.multiply(layer4_1_1, 0.01)

    layer7_1_1 = tf.layers.conv2d(vgg_layer7_out, num_classes,
        kernel_size=1, padding="same", strides=1,
        kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))


    upsampled_2_conv7 = tf.layers.conv2d_transpose(layer7_1_1, num_classes,
        kernel_size=4, strides=2, padding="same",
        kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001),
        kernel_initializer=tf.contrib.layers.xavier_initializer())
    upsampled_2_conv7_with_skipped_connection = tf.add(upsampled_2_conv7, layer4_1_1)

    upsampled_4_conv7 = tf.layers.conv2d_transpose(upsampled_2_conv7_with_skipped_connection,
        num_classes, kernel_size=4, strides=2, padding="same",
        kernel_initializer=tf.contrib.layers.xavier_initializer(),
        kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))
    upsampled_4_conv7_with_skipped_connection = tf.add(upsampled_4_conv7, layer3_1_1)

    upsampled_8_conv7 = tf.layers.conv2d_transpose(upsampled_4_conv7_with_skipped_connection,
        num_classes, kernel_size=16, strides=8, padding="same",
        kernel_initializer=tf.contrib.layers.xavier_initializer(),
        kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))

    return upsampled_8_conv7
tests.test_layers(layers)


def optimize(nn_last_layer, correct_label, learning_rate, num_classes):
    """
    Build the TensorFLow loss and optimizer operations.
    :param nn_last_layer: TF Tensor of the last layer in the neural network
    :param correct_label: TF Placeholder for the correct label image
    :param learning_rate: TF Placeholder for the learning rate
    :param num_classes: Number of classes to classify
    :return: Tuple of (logits, train_op, cross_entropy_loss)
    """
    logits = tf.reshape(nn_last_layer, (-1, num_classes))
    cross_entropy_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
        logits=logits, labels=correct_label))
    reg_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
    cross_entropy_loss += tf.reduce_sum(reg_losses)
    train_op = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cross_entropy_loss)
    return logits, train_op, cross_entropy_loss
tests.test_optimize(optimize)


def train_nn(sess, epochs, batch_size, get_batches_fn, train_op, cross_entropy_loss, input_image,
             correct_label, keep_prob, learning_rate, base_epoch=0):
    """
    Train neural network and print out the loss during training.
    :param sess: TF Session
    :param epochs: int, Number of epochs
    :param batch_size: int, Batch size
    :param get_batches_fn: Function to get batches of training data.  Call using get_batches_fn(batch_size)
    :param train_op: TF Operation to train the neural network
    :param cross_entropy_loss: TF Tensor for the amount of loss
    :param input_image: TF Placeholder for input images
    :param correct_label: TF Placeholder for label images
    :param keep_prob: TF Placeholder for dropout keep probability
    :param learning_rate: TF Placeholder for learning rate
    :param base_epoch: in the case when we need to save checkpoints, we call this function in multi iterations.
        base_epoch is the aboslute number of epochs we have trained so far.
    :return: none
    """
    for i in range(epochs):
        print("Epoch " + str(i + base_epoch))
        total_loss = 0
        for images, labels in get_batches_fn(batch_size):
            _, loss = sess.run([train_op, cross_entropy_loss], feed_dict={
                input_image: images, correct_label: labels,
                keep_prob: 0.5, learning_rate: 0.0001})
            total_loss += loss
        print("Loss: ", total_loss)
tests.test_train_nn(train_nn)


def run():
    num_classes = 2
    image_shape = (160, 576)
    data_dir = './data'
    # image_data_dir = '/small_data'
    image_data_dir = '/data_road'
    runs_dir = './runs'
    if image_data_dir == '/data_road':
        tests.test_for_kitti_dataset(data_dir + image_data_dir)
    save_dir = helper.get_save_dir_name(runs_dir)

    # Download pretrained vgg model
    helper.maybe_download_pretrained_vgg(data_dir)

    # OPTIONAL: Train and Inference on the cityscapes dataset instead of the Kitti dataset.
    # You'll need a GPU with at least 10 teraFLOPS to train on.
    # https://www.cityscapes-dataset.com/

    epochs = 25
    batch_size = 4

    with tf.Session() as sess:
        # Path to vgg model
        vgg_path = os.path.join(data_dir, 'vgg')
        # Create function to get batches
        get_batches_fn = helper.gen_batch_function(os.path.join(data_dir + image_data_dir, 'training'), image_shape)

        # OPTIONAL: Augment Images for better results
        #  https://datascience.stackexchange.com/questions/5224/how-to-prepare-augment-images-for-neural-network

        # Build NN using load_vgg, layers, and optimize function
        image_input, keep_prob, layer3_out, layer4_out, layer7_out = load_vgg(sess, vgg_path)
        nn_last_layer = layers(layer3_out, layer4_out, layer7_out, num_classes)
        learning_rate = tf.placeholder(tf.float32)
        correct_label = tf.placeholder(tf.float32, (None, None, None, num_classes))
        logits, train_op, cross_entropy_loss = optimize(nn_last_layer, correct_label, learning_rate, num_classes)

        sess.run(tf.global_variables_initializer())

        # Train NN using the train_nn function

        # save results in checkpoints, for hyper parameter tuning during training
        if False:
            checkpoints = []
            for i in range(epochs):
                if i >= 20 and i % 5 == 0:
                    checkpoints.append(i)
            if epochs - 1 != i:
                checkpoints.append(epochs - 1)  # add the last epoch

            prev_checkpoint = 0
            for ckpt in checkpoints:
                ckpt_epoch = ckpt - prev_checkpoint
                train_nn(sess, ckpt_epoch, batch_size, get_batches_fn, train_op, cross_entropy_loss, image_input,
                     correct_label, keep_prob, learning_rate, base_epoch=prev_checkpoint)

                # Save inference data using helper.save_inference_samples
                epoch_dir = "/epoch" + str(ckpt)
                helper.save_inference_samples(save_dir + epoch_dir, data_dir + image_data_dir, sess,
                    image_shape, logits, keep_prob, image_input)
                prev_checkpoint = ckpt

        # train in one go without saving checkpoints
        else:
            train_nn(sess, epochs, batch_size, get_batches_fn, train_op, cross_entropy_loss, image_input,
                correct_label, keep_prob, learning_rate)
            helper.save_inference_samples(save_dir, data_dir + image_data_dir, sess,
                image_shape, logits, keep_prob, image_input)

if __name__ == '__main__':
    run()

#!/usr/bin/env python3
import os.path
import tensorflow as tf
import helper
import warnings
from distutils.version import LooseVersion
import project_tests as tests


# Check TensorFlow Version
assert LooseVersion(tf.__version__) >= LooseVersion('1.0'), 'Please use TensorFlow version 1.0 or newer.  You are using {}'.format(tf.__version__)
print('TensorFlow Version: {}'.format(tf.__version__))

# Check for a GPU
if not tf.test.gpu_device_name():
    warnings.warn('No GPU found. Please use a GPU to train your neural network.')
else:
    print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))


def load_vgg(sess, vgg_path):
    """
    Load Pretrained VGG Model into TensorFlow.
    :param sess: TensorFlow Session
    :param vgg_path: Path to vgg folder, containing "variables/" and "saved_model.pb"
    :return: Tuple of Tensors from VGG model (image_input, keep_prob, layer3_out, layer4_out, layer7_out)
    """
    tf.saved_model.loader.load(sess, ['vgg16'], vgg_path)
    graph = tf.get_default_graph()
    image_input = graph.get_tensor_by_name('image_input:0')
    keep_prob = graph.get_tensor_by_name('keep_prob:0')
    layer3_out = graph.get_tensor_by_name('layer3_out:0')
    layer4_out = graph.get_tensor_by_name('layer4_out:0')
    layer7_out = graph.get_tensor_by_name('layer7_out:0')

    return image_input, keep_prob, layer3_out, layer4_out, layer7_out
tests.test_load_vgg(load_vgg, tf)


def layers(vgg_layer3_out, vgg_layer4_out, vgg_layer7_out, num_classes):
    """
    Create the layers for a fully convolutional network.  Build skip-layers using the vgg layers.
    :param vgg_layer3_out: TF Tensor for VGG Layer 3 output
    :param vgg_layer4_out: TF Tensor for VGG Layer 4 output
    :param vgg_layer7_out: TF Tensor for VGG Layer 7 output
    :param num_classes: Number of classes to classify
    :return: The Tensor for the last layer of output
    """
    # 1 by 1 convolution
    layer3_1_1 = tf.layers.conv2d(vgg_layer3_out, num_classes,
        kernel_size=1, padding="same", strides=1,
        kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))
    # scaling
    layer3_1_1 = tf.multiply(layer3_1_1, 0.0001)

    layer4_1_1 = tf.layers.conv2d(vgg_layer4_out, num_classes,
        kernel_size=1, padding="same", strides=1,
        kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))
    # scaling
    layer4_1_1 = tf.multiply(layer4_1_1, 0.01)

    layer7_1_1 = tf.layers.conv2d(vgg_layer7_out, num_classes,
        kernel_size=1, padding="same", strides=1,
        kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))


    upsampled_2_conv7 = tf.layers.conv2d_transpose(layer7_1_1, num_classes,
        kernel_size=4, strides=2, padding="same",
        kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001),
        kernel_initializer=tf.contrib.layers.xavier_initializer())
    upsampled_2_conv7_with_skipped_connection = tf.add(upsampled_2_conv7, layer4_1_1)

    upsampled_4_conv7 = tf.layers.conv2d_transpose(upsampled_2_conv7_with_skipped_connection,
        num_classes, kernel_size=4, strides=2, padding="same",
        kernel_initializer=tf.contrib.layers.xavier_initializer(),
        kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))
    upsampled_4_conv7_with_skipped_connection = tf.add(upsampled_4_conv7, layer3_1_1)

    upsampled_8_conv7 = tf.layers.conv2d_transpose(upsampled_4_conv7_with_skipped_connection,
        num_classes, kernel_size=16, strides=8, padding="same",
        kernel_initializer=tf.contrib.layers.xavier_initializer(),
        kernel_regularizer=tf.contrib.layers.l2_regularizer(0.001))

    return upsampled_8_conv7
tests.test_layers(layers)


def optimize(nn_last_layer, correct_label, learning_rate, num_classes):
    """
    Build the TensorFLow loss and optimizer operations.
    :param nn_last_layer: TF Tensor of the last layer in the neural network
    :param correct_label: TF Placeholder for the correct label image
    :param learning_rate: TF Placeholder for the learning rate
    :param num_classes: Number of classes to classify
    :return: Tuple of (logits, train_op, cross_entropy_loss)
    """
    logits = tf.reshape(nn_last_layer, (-1, num_classes))
    cross_entropy_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
        logits=logits, labels=correct_label))
    reg_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
    cross_entropy_loss += tf.reduce_sum(reg_losses)
    train_op = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cross_entropy_loss)
    return logits, train_op, cross_entropy_loss
tests.test_optimize(optimize)


def train_nn(sess, epochs, batch_size, get_batches_fn, train_op, cross_entropy_loss, input_image,
             correct_label, keep_prob, learning_rate, base_epoch=0):
    """
    Train neural network and print out the loss during training.
    :param sess: TF Session
    :param epochs: int, Number of epochs
    :param batch_size: int, Batch size
    :param get_batches_fn: Function to get batches of training data.  Call using get_batches_fn(batch_size)
    :param train_op: TF Operation to train the neural network
    :param cross_entropy_loss: TF Tensor for the amount of loss
    :param input_image: TF Placeholder for input images
    :param correct_label: TF Placeholder for label images
    :param keep_prob: TF Placeholder for dropout keep probability
    :param learning_rate: TF Placeholder for learning rate
    :param base_epoch: in the case when we need to save checkpoints, we call this function in multi iterations.
        base_epoch is the aboslute number of epochs we have trained so far.
    :return: none
    """
    for i in range(epochs):
        print("Epoch " + str(i + base_epoch))
        total_loss = 0
        for images, labels in get_batches_fn(batch_size):
            _, loss = sess.run([train_op, cross_entropy_loss], feed_dict={
                input_image: images, correct_label: labels,
                keep_prob: 0.5, learning_rate: 0.0001})
            total_loss += loss
        print("Loss: ", total_loss)
tests.test_train_nn(train_nn)


def run():
    num_classes = 2
    image_shape = (160, 576)
    data_dir = './data'
    # image_data_dir = '/small_data'
    image_data_dir = '/data_road'
    runs_dir = './runs'
    if image_data_dir == '/data_road':
        tests.test_for_kitti_dataset(data_dir + image_data_dir)
    save_dir = helper.get_save_dir_name(runs_dir)

    # Download pretrained vgg model
    helper.maybe_download_pretrained_vgg(data_dir)

    # OPTIONAL: Train and Inference on the cityscapes dataset instead of the Kitti dataset.
    # You'll need a GPU with at least 10 teraFLOPS to train on.
    # https://www.cityscapes-dataset.com/

    epochs = 25
    batch_size = 4

    with tf.Session() as sess:
        # Path to vgg model
        vgg_path = os.path.join(data_dir, 'vgg')
        # Create function to get batches
        get_batches_fn = helper.gen_batch_function(os.path.join(data_dir + image_data_dir, 'training'), image_shape)

        # OPTIONAL: Augment Images for better results
        #  https://datascience.stackexchange.com/questions/5224/how-to-prepare-augment-images-for-neural-network

        # Build NN using load_vgg, layers, and optimize function
        image_input, keep_prob, layer3_out, layer4_out, layer7_out = load_vgg(sess, vgg_path)
        nn_last_layer = layers(layer3_out, layer4_out, layer7_out, num_classes)
        learning_rate = tf.placeholder(tf.float32)
        correct_label = tf.placeholder(tf.float32, (None, None, None, num_classes))
        logits, train_op, cross_entropy_loss = optimize(nn_last_layer, correct_label, learning_rate, num_classes)

        sess.run(tf.global_variables_initializer())

        # Train NN using the train_nn function

        # save results in checkpoints, for hyper parameter tuning during training
        if False:
            checkpoints = []
            for i in range(epochs):
                if i >= 20 and i % 5 == 0:
                    checkpoints.append(i)
            if epochs - 1 != i:
                checkpoints.append(epochs - 1)  # add the last epoch

            prev_checkpoint = 0
            for ckpt in checkpoints:
                ckpt_epoch = ckpt - prev_checkpoint
                train_nn(sess, ckpt_epoch, batch_size, get_batches_fn, train_op, cross_entropy_loss, image_input,
                     correct_label, keep_prob, learning_rate, base_epoch=prev_checkpoint)

                # Save inference data using helper.save_inference_samples
                epoch_dir = "/epoch" + str(ckpt)
                helper.save_inference_samples(save_dir + epoch_dir, data_dir + image_data_dir, sess,
                    image_shape, logits, keep_prob, image_input)
                prev_checkpoint = ckpt

        # train in one go without saving checkpoints
        else:
            train_nn(sess, epochs, batch_size, get_batches_fn, train_op, cross_entropy_loss, image_input,
                correct_label, keep_prob, learning_rate)
            helper.save_inference_samples(save_dir, data_dir + image_data_dir, sess,
                image_shape, logits, keep_prob, image_input)

if __name__ == '__main__':
    run()

#!/usr/bin/env python3 # Check TensorFlow Version # Check for a GPU Load Pretrained VGG Model into TensorFlow.     :param sess: TensorFlow Session     :param vgg_path: Path to vgg folder, containing "variables/" and "saved_model.pb"     :return: Tuple of Tensors from VGG model (image_input, keep_prob, layer3_out, layer4_out, layer7_out) Create the layers for a fully convolutional network.  Build skip-layers using the vgg layers.     :param vgg_layer3_out: TF Tensor for VGG Layer 3 output     :param vgg_layer4_out: TF Tensor for VGG Layer 4 output     :param vgg_layer7_out: TF Tensor for VGG Layer 7 output     :param num_classes: Number of classes to classify     :return: The Tensor for the last layer of output # 1 by 1 convolution # scaling # scaling Build the TensorFLow loss and optimizer operations.     :param nn_last_layer: TF Tensor of the last layer in the neural network     :param correct_label: TF Placeholder for the correct label image     :param learning_rate: TF Placeholder for the learning rate     :param num_classes: Number of classes to classify     :return: Tuple of (logits, train_op, cross_entropy_loss) Train neural network and print out the loss during training.     :param sess: TF Session     :param epochs: int, Number of epochs     :param batch_size: int, Batch size     :param get_batches_fn: Function to get batches of training data.  Call using get_batches_fn(batch_size)     :param train_op: TF Operation to train the neural network     :param cross_entropy_loss: TF Tensor for the amount of loss     :param input_image: TF Placeholder for input images     :param correct_label: TF Placeholder for label images     :param keep_prob: TF Placeholder for dropout keep probability     :param learning_rate: TF Placeholder for learning rate     :param base_epoch: in the case when we need to save checkpoints, we call this function in multi iterations.         base_epoch is the aboslute number of epochs we have trained so far.     :return: none # image_data_dir = '/small_data' # Download pretrained vgg model # OPTIONAL: Train and Inference on the cityscapes dataset instead of the Kitti dataset. # You'll need a GPU with at least 10 teraFLOPS to train on. # https://www.cityscapes-dataset.com/ # Path to vgg model # Create function to get batches # OPTIONAL: Augment Images for better results #  https://datascience.stackexchange.com/questions/5224/how-to-prepare-augment-images-for-neural-network # Build NN using load_vgg, layers, and optimize function # Train NN using the train_nn function # save results in checkpoints, for hyper parameter tuning during training # add the last epoch # Save inference data using helper.save_inference_samples # train in one go without saving checkpoints