# -*- coding: utf-8 -*-
"""
Created on Tue Jun 26 16:34:21 2018

@author: LiHongWang
"""

import os
import tensorflow as tf
from model import fcn_vgg
from model import fcn_mobile
from model import fcn_resnet_v2
from data import input_data


slim = tf.contrib.slim



def main():
    
    num_classes=2
    tfRecorf_dir= 'D:/dataSet/kitti/road/sub_um_lane_tra66.tfrecord'
    
    train_dir = './fm2/'   
    if not os.path.exists(train_dir):
        os.makedirs(train_dir)
        
    with tf.Graph().as_default():
        global_step = tf.contrib.framework.get_or_create_global_step()
        tf.logging.set_verbosity(tf.logging.INFO)
        
        with tf.device("/cpu:0"):
            samples=input_data.get_images_labels(tfRecorf_dir,num_classes,66,
                                                 crop_size=[224,224],
                                                 batch_size=4)
        
              

            batch_queue = slim.prefetch_queue.prefetch_queue(samples,
                                                             capacity=128 )
 
        
            tra_batch = batch_queue.dequeue()

        logit,prediction=fcn_mobile.fcn_mobv1(tra_batch['image'],num_classes)
#        logit,prediction=fcn_vgg.fcn_vgg16(tra_batch['image'],num_classes)
        
#        logit,prediction=fcn_resnet_v2.fcn_res101(tra_batch['image'],num_classes)
        
        
        
        cross_entropy=tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logit,
                           labels=tf.squeeze(tra_batch['label'], squeeze_dims=[3]),name="entropy")
    

    
    
        loss = tf.reduce_mean(cross_entropy,name='loss')
        slim.losses.add_loss(loss) 
        
        total_loss = slim.losses.get_total_loss()


#        print("image", tra_batch['image'])
#        print("label", tf.cast(tra_batch['label']*255, tf.uint8))
#        print("prediction", tf.cast(prediction*255, tf.uint8))

    # Create some summaries to visualize the training process:
        tf.summary.scalar('losses/Total_Loss', total_loss)
        tf.summary.image("image", tra_batch['image'], max_outputs=4)
        tf.summary.image("label", tf.cast(tra_batch['label']*255, tf.uint8), max_outputs=4)
        tf.summary.image("prediction", tf.cast(prediction*255, tf.uint8), max_outputs=4)
        
        
        
  
        lr = tf.train.exponential_decay(0.001,
                                  global_step,
                                  10000,
                                  0.8,
                                  staircase=True)

        #lr = tf.constant(0.001, tf.float32)
        tf.summary.scalar('learning_rate', lr)
        for variable in slim.get_model_variables():
           tf.summary.histogram(variable.op.name, variable)
    
    
    # Specify the optimizer and create the train op:
        optimizer = tf.train.RMSPropOptimizer(lr,0.9)
        train_op = slim.learning.create_train_op(total_loss, optimizer)
    
    # Run the training:
    
        
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.7)

        config=tf.ConfigProto(gpu_options=gpu_options)
        final_loss = slim.learning.train(train_op,
                                     logdir=train_dir,
                                     log_every_n_steps=100,
                                     save_summaries_secs=20,
                                     save_interval_secs=1800,
                                     init_fn=None,#fcn_mobile.get_init_fn(),
                                     session_config=config,
                                     number_of_steps=65000)
          
    print('Finished training. Last batch loss %f' % final_loss)
    
if __name__=='__main__':
    main()
[((26, 11, 26, 36), 'os.path.exists', 'os.path.exists', ({(26, 26, 26, 35): 'train_dir'}, {}), '(train_dir)', False, 'import os\n'), ((27, 8, 27, 30), 'os.makedirs', 'os.makedirs', ({(27, 20, 27, 29): 'train_dir'}, {}), '(train_dir)', False, 'import os\n'), ((30, 22, 30, 70), 'tensorflow.contrib.framework.get_or_create_global_step', 'tf.contrib.framework.get_or_create_global_step', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((31, 8, 31, 49), 'tensorflow.logging.set_verbosity', 'tf.logging.set_verbosity', ({(31, 33, 31, 48): 'tf.logging.INFO'}, {}), '(tf.logging.INFO)', True, 'import tensorflow as tf\n'), ((46, 25, 46, 77), 'model.fcn_mobile.fcn_mobv1', 'fcn_mobile.fcn_mobv1', ({(46, 46, 46, 64): "tra_batch['image']", (46, 65, 46, 76): 'num_classes'}, {}), "(tra_batch['image'], num_classes)", False, 'from model import fcn_mobile\n'), ((59, 15, 59, 56), 'tensorflow.reduce_mean', 'tf.reduce_mean', (), '', True, 'import tensorflow as tf\n'), ((70, 8, 70, 58), 'tensorflow.summary.scalar', 'tf.summary.scalar', ({(70, 26, 70, 45): '"""losses/Total_Loss"""', (70, 47, 70, 57): 'total_loss'}, {}), "('losses/Total_Loss', total_loss)", True, 'import tensorflow as tf\n'), ((71, 8, 71, 68), 'tensorflow.summary.image', 'tf.summary.image', (), '', True, 'import tensorflow as tf\n'), ((78, 13, 82, 49), 'tensorflow.train.exponential_decay', 'tf.train.exponential_decay', (), '', True, 'import tensorflow as tf\n'), ((85, 8, 85, 46), 'tensorflow.summary.scalar', 'tf.summary.scalar', ({(85, 26, 85, 41): '"""learning_rate"""', (85, 43, 85, 45): 'lr'}, {}), "('learning_rate', lr)", True, 'import tensorflow as tf\n'), ((91, 20, 91, 53), 'tensorflow.train.RMSPropOptimizer', 'tf.train.RMSPropOptimizer', ({(91, 46, 91, 48): 'lr', (91, 49, 91, 52): '0.9'}, {}), '(lr, 0.9)', True, 'import tensorflow as tf\n'), ((97, 22, 97, 72), 'tensorflow.GPUOptions', 'tf.GPUOptions', (), '', True, 'import tensorflow as tf\n'), ((99, 15, 99, 54), 'tensorflow.ConfigProto', 'tf.ConfigProto', (), '', True, 'import tensorflow as tf\n'), ((33, 13, 33, 32), 'tensorflow.device', 'tf.device', ({(33, 23, 33, 31): '"""/cpu:0"""'}, {}), "('/cpu:0')", True, 'import tensorflow as tf\n'), ((34, 20, 36, 62), 'data.input_data.get_images_labels', 'input_data.get_images_labels', (), '', False, 'from data import input_data\n'), ((72, 34, 72, 75), 'tensorflow.cast', 'tf.cast', ({(72, 42, 72, 64): "(tra_batch['label'] * 255)", (72, 66, 72, 74): 'tf.uint8'}, {}), "(tra_batch['label'] * 255, tf.uint8)", True, 'import tensorflow as tf\n'), ((73, 39, 73, 72), 'tensorflow.cast', 'tf.cast', ({(73, 47, 73, 61): '(prediction * 255)', (73, 63, 73, 71): 'tf.uint8'}, {}), '(prediction * 255, tf.uint8)', True, 'import tensorflow as tf\n'), ((87, 11, 87, 59), 'tensorflow.summary.histogram', 'tf.summary.histogram', ({(87, 32, 87, 48): 'variable.op.name', (87, 50, 87, 58): 'variable'}, {}), '(variable.op.name, variable)', True, 'import tensorflow as tf\n'), ((29, 9, 29, 19), 'tensorflow.Graph', 'tf.Graph', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((54, 34, 54, 82), 'tensorflow.squeeze', 'tf.squeeze', (), '', True, 'import tensorflow as tf\n')]