"""

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)

"""

# Use tf.saved_model.loader.load to load the model and weights

vgg_tag = 'vgg16'

vgg_input_tensor_name = 'image_input:0'

vgg_keep_prob_tensor_name = 'keep_prob:0'

vgg_layer3_out_tensor_name = 'layer3_out:0'

vgg_layer4_out_tensor_name = 'layer4_out:0'

vgg_layer7_out_tensor_name = 'layer7_out:0'

tf.saved_model.loader.load(sess, [vgg_tag], vgg_path)

graph = tf.get_default_graph()

t_input = graph.get_tensor_by_name(vgg_input_tensor_name)

t_keep_prob = graph.get_tensor_by_name(vgg_keep_prob_tensor_name)

t_layer3_out = graph.get_tensor_by_name(vgg_layer3_out_tensor_name)

t_layer4_out = graph.get_tensor_by_name(vgg_layer4_out_tensor_name)

t_layer7_out = graph.get_tensor_by_name(vgg_layer7_out_tensor_name)

"""

Create the layers for a fully convolutional network. Build skip-layers using the vgg layers.

:param vgg_layer7_out: TF Tensor for VGG Layer 3 output

:param vgg_layer4_out: TF Tensor for VGG Layer 4 output

:param vgg_layer3_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

"""

# 1x1 convolution layer for feature extraction

tf_conv1x1 = tf.layers.conv2d(vgg_layer7_out, num_classes, 1, 1,

kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3),

kernel_initializer=tf.truncated_normal_initializer(stddev=0.01))

# conv2d-transpose for 2x upsampling

tf_2x = tf.layers.conv2d_transpose(tf_conv1x1, num_classes, 4, 2, padding='SAME',

kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3),

kernel_initializer=tf.truncated_normal_initializer(stddev=0.01))

# combine with pooling layer 4

tf_skip1 = tf.add(tf_2x, tf.layers.conv2d(vgg_layer4_out, num_classes, 1, 1,

kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3),

kernel_initializer=tf.truncated_normal_initializer(stddev=0.01)))

# perform conv2d-transpose for 2x upsampling again.

# output: 4x features + 2x pool4

tf_4x = tf.layers.conv2d_transpose(tf_skip1, num_classes, 4, 2, padding='SAME',

kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3),

kernel_initializer=tf.truncated_normal_initializer(stddev=0.01))

# combile with pooling layer 3

tf_skip2 = tf.add(tf_4x, tf.layers.conv2d(vgg_layer3_out, num_classes, 1, 1,

kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3),

kernel_initializer=tf.truncated_normal_initializer(stddev=0.01)))

# perform conv2d-transpose for 2x upsampling again.

# output: 8x features + 4x pool4 + 2x pool3

tf_final = tf.layers.conv2d_transpose(tf_skip2, num_classes, 16, 8, padding='SAME',

kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-3),

kernel_initializer=tf.truncated_normal_initializer(stddev=0.01))

tf.stop_gradient(vgg_layer3_out)

tf.stop_gradient(vgg_layer4_out)

tf.stop_gradient(vgg_layer7_out)

"""

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

labels = tf.reshape(correct_label, (-1, num_classes))

cross_entropy_loss = tf.reduce_mean(

tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=labels))

# cost without l2 regularizer

total_loss = cross_entropy_loss

# cost with l2 regularizer

# l2_reg = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)

# l2_loss = tf.reduce_sum(l2_reg)

# total_loss = cross_entropy_loss + l2_loss

train_op = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(total_loss)

correct_label, keep_prob, learning_rate):

"""

Train neural network and print out the loss during training.

:param sess: TF Session

:param epochs: Number of epochs

:param batch_size: 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

"""

params = {

'learning_rate': 0.0001,

'keep_prob': 0.5

}

sess.run(tf.global_variables_initializer())

sess.run(tf.local_variables_initializer())

for e in range(epochs):

# size of training set: 290

total_lost = 0

processed_samples = 0

num_batches = math.ceil(290 / batch_size)

loop = tqdm(get_batches_fn(batch_size), total=num_batches)

for input_images, labels in loop:

# change labels to 1.0 with True and 0.0 with False

labels = labels.astype(np.float32)

_, loss = sess.run([train_op, cross_entropy_loss],

feed_dict={

input_image: input_images,

correct_label: labels,

keep_prob: params['keep_prob'],

learning_rate: params['learning_rate']

})

num_samples = input_images.shape[0]

processed_samples += num_samples

total_lost += (loss * num_samples)

mean_lost = total_lost / processed_samples

num_classes = 2

image_shape = (160, 576)

data_dir = './data'

runs_dir = './runs'

tests.test_for_kitti_dataset(data_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 = 30

batch_size = 5

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, 'data_road/training'), image_shape)

# OPTIONAL: Augment Images for better results

# https://datascience.stackexchange.com/questions/5224/how-to-prepare-augment-images-for-neural-network

# placeholder for learning rate, input_image and correct_label

t_learning_rate = tf.placeholder(tf.float32,

name='learning_rate')

t_correct_label = tf.placeholder(tf.float32,

shape=(None, image_shape[0], image_shape[1], num_classes),

name='correct_label')

t_input, t_keep_prob, t_layer3_out, t_layer4_out, t_layer7_out = load_vgg(sess, vgg_path)

output = layers(t_layer3_out, t_layer4_out, t_layer7_out, num_classes)

logits, train_op, loss = optimize(output, t_correct_label, t_learning_rate, num_classes)

train_nn(sess, epochs, batch_size, get_batches_fn, train_op, loss, t_input,

t_correct_label, t_keep_prob, t_learning_rate)

helper.save_inference_samples(runs_dir, data_dir, sess, image_shape, logits, t_keep_prob, t_input)