text: """Run DeepLab-ResNet on a given image.
  
  This script computes a segmentation mask for a given image.
  """
  
  from __future__ import print_function
  
  import argparse
  from datetime import datetime
  import os
  import sys
  import time
  
  from PIL import Image
  
  import tensorflow as tf
  import numpy as np
  
  from deeplab_resnet import DeepLabResNetModel, ImageReader, decode_labels, prepare_label
  
  IMG_MEAN = np.array((104.00698793,116.66876762,122.67891434), dtype=np.float32)
      
  NUM_CLASSES = 21
  SAVE_DIR = 'output/'
  RESTORE_FROM = 'snapshots_finetune/model.ckpt-19900'
  INPUT_SIZE=(512,512)
  
  def get_arguments():
      """Parse all the arguments provided from the CLI.
      
      Returns:
        A list of parsed arguments.
      """
      parser = argparse.ArgumentParser(description="DeepLabLFOV Network Inference.")
      parser.add_argument("--img_path", type=str,
                          help="Path to the RGB image file.")
      parser.add_argument("--model_weights", type=str,default = RESTORE_FROM,
                          help="Path to the file with model weights.")
      parser.add_argument("--num-classes", type=int, default=NUM_CLASSES,
                          help="Number of classes to predict (including background).")
      parser.add_argument("--save-dir", type=str, default=SAVE_DIR,
                          help="Where to save predicted mask.")
      return parser.parse_args()
  
  def load(saver, sess, ckpt_path):
      '''Load trained weights.
      
      Args:
        saver: TensorFlow saver object.
        sess: TensorFlow session.
        ckpt_path: path to checkpoint file with parameters.
      ''' 
      saver.restore(sess, ckpt_path)
      print("Restored model parameters from {}".format(ckpt_path))
  
  def main():
      """Create the model and start the evaluation process."""
      args = get_arguments()
      
      # Prepare image.
      for i in os.listdir(args.img_path):
          name = i.split('.')
          img = tf.image.decode_jpeg(tf.read_file(os.path.join(args.img_path,i)), channels=3)
          img_shape = img.get_shape().as_list()
          # Convert RGB to BGR.
          img_r, img_g, img_b = tf.split(axis=2, num_or_size_splits=3, value=img)
          img = tf.cast(tf.concat(axis=2, values=[img_b, img_g, img_r]), dtype=tf.float32)
          # Extract mean.
          img -= IMG_MEAN 
          h,w = INPUT_SIZE
          
          if img_shape[0]!=h and img_shape[1]!=w:
              img = tf.image.resize_images(img,[h,w])
          
          # Create network.
          net = DeepLabResNetModel({'data': tf.expand_dims(img, dim=0)}, is_training=False, num_classes=args.num_classes)
      
          # Which variables to load.
          restore_var = tf.global_variables()
      
          # Predictions.
          raw_output = net.layers['fc1_voc12']
          raw_output_up = tf.image.resize_bilinear(raw_output, tf.shape(img)[0:2,])
          raw_output_up = tf.argmax(raw_output_up, dimension=3)
          pred = tf.expand_dims(raw_output_up, dim=3)
      
          
          # Set up TF session and initialize variables. 
          config = tf.ConfigProto()
          config.gpu_options.allow_growth = True
          sess = tf.Session(config=config)
          init = tf.global_variables_initializer()
          
          sess.run(init)
          
          # Load weights.
          loader = tf.train.Saver(var_list=restore_var)
          load(loader, sess, args.model_weights)
          
          # Perform inference.
          start = time.time()
          preds = sess.run(pred)
          end = time.time()
          
          print('Inference single image need: {:.5f}'.format(end-start))
          
          msk = decode_labels(preds, num_classes=args.num_classes)
          im = Image.fromarray(msk[0])
          im1 = im.resize((512,512),Image.BICUBIC)
          if not os.path.exists(args.save_dir):
              os.makedirs(args.save_dir)
  #        im.save(args.save_dir + 'mask.png')
          im1.save(args.save_dir + name[0]+'_predict.png')
          tf.reset_default_graph()
  #        print('The output file has been saved to {}'.format(args.save_dir + 'mask.png'))
  
      
  if __name__ == '__main__':
      main()
  