# -*- 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, 'os.path.exists', 'os.path.exists', (['train_dir'], {}), False, 'import os\n'), (27, 'os.makedirs', 'os.makedirs', (['train_dir'], {}), False, 'import os\n'), (30, 'tensorflow.contrib.framework.get_or_create_global_step', 'tf.contrib.framework.get_or_create_global_step', ([], {}), True, 'import tensorflow as tf\n'), (31, 'tensorflow.logging.set_verbosity', 'tf.logging.set_verbosity', (['tf.logging.INFO'], {}), True, 'import tensorflow as tf\n'), (46, 'model.fcn_mobile.fcn_mobv1', 'fcn_mobile.fcn_mobv1', (["tra_batch['image']", 'num_classes'], {}), False, 'from model import fcn_mobile\n'), (59, 'tensorflow.reduce_mean', 'tf.reduce_mean', (['cross_entropy'], {'name': '"""loss"""'}), True, 'import tensorflow as tf\n'), (70, 'tensorflow.summary.scalar', 'tf.summary.scalar', (['"""losses/Total_Loss"""', 'total_loss'], {}), True, 'import tensorflow as tf\n'), (71, 'tensorflow.summary.image', 'tf.summary.image', (['"""image"""', "tra_batch['image']"], {'max_outputs': '(4)'}), True, 'import tensorflow as tf\n'), (78, 'tensorflow.train.exponential_decay', 'tf.train.exponential_decay', (['(0.001)', 'global_step', '(10000)', '(0.8)'], {'staircase': '(True)'}), True, 'import tensorflow as tf\n'), (85, 'tensorflow.summary.scalar', 'tf.summary.scalar', (['"""learning_rate"""', 'lr'], {}), True, 'import tensorflow as tf\n'), (91, 'tensorflow.train.RMSPropOptimizer', 'tf.train.RMSPropOptimizer', (['lr', '(0.9)'], {}), True, 'import tensorflow as tf\n'), (97, 'tensorflow.GPUOptions', 'tf.GPUOptions', ([], {'per_process_gpu_memory_fraction': '(0.7)'}), True, 'import tensorflow as tf\n'), (99, 'tensorflow.ConfigProto', 'tf.ConfigProto', ([], {'gpu_options': 'gpu_options'}), True, 'import tensorflow as tf\n'), (33, 'tensorflow.device', 'tf.device', (['"""/cpu:0"""'], {}), True, 'import tensorflow as tf\n'), (34, 'data.input_data.get_images_labels', 'input_data.get_images_labels', (['tfRecorf_dir', 'num_classes', '(66)'], {'crop_size': '[224, 224]', 'batch_size': '(4)'}), False, 'from data import input_data\n'), (72, 'tensorflow.cast', 'tf.cast', (["(tra_batch['label'] * 255)", 'tf.uint8'], {}), True, 'import tensorflow as tf\n'), (73, 'tensorflow.cast', 'tf.cast', (['(prediction * 255)', 'tf.uint8'], {}), True, 'import tensorflow as tf\n'), (87, 'tensorflow.summary.histogram', 'tf.summary.histogram', (['variable.op.name', 'variable'], {}), True, 'import tensorflow as tf\n'), (29, 'tensorflow.Graph', 'tf.Graph', ([], {}), True, 'import tensorflow as tf\n'), (54, 'tensorflow.squeeze', 'tf.squeeze', (["tra_batch['label']"], {'squeeze_dims': '[3]'}), True, 'import tensorflow as tf\n')]