text: """
   @Time    : 201/28/19 09:07
   @Author  : TaylorMei
   @Email   : mhy845879017@gmail.com
   
   @Project : iccv
   @File    : infer_mhy.py
   @Function:
   
  """
  import numpy as np
  import os
  import time
  
  import torch
  from PIL import Image
  from torch.autograd import Variable
  from torchvision import transforms
  import skimage.io
  
  from config import msd_testing_root
  from misc import check_mkdir, crf_refine
  from model.our2 import OUR2
  
  device_ids = [0]
  torch.cuda.set_device(device_ids[0])
  
  ckpt_path = './ckpt'
  exp_name = 'OUR2'
  args = {
      'snapshot': '140',
      'scale': 480,
      'crf': True
  }
  
  img_transform = transforms.Compose([
      transforms.Resize((args['scale'], args['scale'])),
      transforms.ToTensor(),
      transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  ])
  
  to_test = {'msd9': msd_testing_root}
  
  to_pil = transforms.ToPILImage()
  
  
  def main():
      net = OUR2().cuda(device_ids[0])
  
      if len(args['snapshot']) > 0:
          print('Load snapshot {} for testing'.format(args['snapshot']))
          net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
          print('Load {} succeed!'.format(os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
  
      net.eval()
      with torch.no_grad():
          for name, root in to_test.items():
              # img_list = [img_name for img_name in os.listdir(os.path.join(root, 'image')) if img_name.endswith('.jpg')]
              img_list = [img_name for img_name in os.listdir(os.path.join(root, 'image')) ]
              start = time.time()
              for idx, img_name in enumerate(img_list):
                  print('predicting for {}: {:>4d} / {}'.format(name, idx + 1, len(img_list)))
                  check_mkdir(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot'])))
                  img = Image.open(os.path.join(root, 'image', img_name))
                  if img.mode != 'RGB':
                      img = img.convert('RGB')
                      print("{} is a gray image.".format(name))
                  w, h = img.size
                  img_var = Variable(img_transform(img).unsqueeze(0)).cuda()
                  _, _, _, _, f = net(img_var)
                  output = f.data.squeeze(0).cpu()
                  prediction = np.array(transforms.Resize((h, w))(to_pil(output)))
                  if args['crf']:
                      prediction = crf_refine(np.array(img.convert('RGB')), prediction)
  
                  # Image.fromarray(prediction).save(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), img_name[:-4] + ".png"))
                  # skimage.io.imsave(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), img_name[:-4] + ".png"), np.where(prediction>=127.5, 255, 0).astype(np.uint8))
                  # skimage.io.imsave(os.path.join(ckpt_path, exp_name, '%s_%s' % (exp_name, args['snapshot']), img_name[:-4] + ".png"), prediction.astype(np.uint8))
                  check_mkdir(os.path.join(msd_testing_root, 'mirror_map'))
                  Image.fromarray(prediction).save(os.path.join(msd_testing_root, 'mirror_map', img_name[:-4] + ".png"))
              end = time.time()
              print("Average Time Is : {:.2f}".format((end - start) / len(img_list)))
  
  
  if __name__ == '__main__':
      main()
  