import tensorflow as tf

    if FLAGS.mode == 'train':
        train_dataset_reader = dataset.BatchDatset(train_records, image_options)
    validation_dataset_reader = dataset.BatchDatset(valid_records, image_options)

    sess = tf.Session()

    print("Setting up Saver...")
    saver = tf.train.Saver()

    # create two summary writers to show training loss and validation loss in the same graph
    # need to create two folders 'train' and 'validation' inside FLAGS.logs_dir
    train_writer = tf.summary.FileWriter(FLAGS.logs_dir + '/train', sess.graph)
    validation_writer = tf.summary.FileWriter(FLAGS.logs_dir + '/validation')

    sess.run(tf.global_variables_initializer())
    ckpt = tf.train.get_checkpoint_state(FLAGS.logs_dir)
    if ckpt and ckpt.model_checkpoint_path:
        saver.restore(sess, ckpt.model_checkpoint_path)
        print("Model restored...")

    if FLAGS.mode == "train":
        for itr in xrange(MAX_ITERATION):
            
            train_images, train_annotations = train_dataset_reader.next_batch(FLAGS.batch_size)
            z_ = np.random.uniform(low=-1.0, high=1.0, size=(FLAGS.batch_size,4,4,128))
         #   print(train_images)
            feed_dict = {image: train_images, annotation: train_annotations, keep_probability: 0.85, z: z_}
           #train_images[:,50:100,50:100,:] =0
            v = 0
            
