text: """
  convert  a face to another person
  
  """
  from models.swapnet import SwapNet
  import torch
  from alfred.dl.torch.common import device
  import cv2
  import numpy as np
  from dataset.training_data import random_warp
  from utils.umeyama import umeyama
  
  mean_value = np.array([0.03321508, 0.05035182, 0.02038819])
  
  
  def process_img(ori_img):
      img = cv2.resize(ori_img, (256, 256))
      range_ = np.linspace( 128-80, 128+80, 5 )
      mapx = np.broadcast_to( range_, (5,5) )
      mapy = mapx.T
  
      # warp image like in the training
      mapx = mapx + np.random.normal( size=(5,5), scale=5 )
      mapy = mapy + np.random.normal( size=(5,5), scale=5 )
      interp_mapx = cv2.resize(mapx, (80, 80))[8:72, 8:72].astype('float32')
      interp_mapy = cv2.resize(mapy, (80, 80))[8:72, 8:72].astype('float32')
      warped_image = cv2.remap(img, interp_mapx, interp_mapy, cv2.INTER_LINEAR)
      return warped_image
  
  
  def load_img():
      a = 'images/34600_test_A_target.png'
      img = cv2.imread(a) / 255.
      return img
  
  
  def predict():
      # convert trump to cage
      # img_f = 'data/trump/51834796.jpg'
      # img_f = 'data/trump/494045244.jpg'
      # NOTE: using face extracted (not original image)
      img_f = 'data/trump/464669134_face_0.png'
  
      ori_img = cv2.imread(img_f)
      img = cv2.resize(ori_img, (64, 64)) / 255.
      img = np.rot90(img)
      # img = load_img()
      in_img = np.array(img, dtype=np.float).transpose(2, 1, 0)
  
      # normalize img
      in_img = torch.Tensor(in_img).to(device).unsqueeze(0)
      model = SwapNet().to(device)
      if torch.cuda.is_available():
          checkpoint = torch.load('checkpoint/faceswap_trump_cage_64x64.pth')
      else:
          checkpoint = torch.load('checkpoint/faceswap_trump_cage_64x64.pth', map_location={'cuda:0': 'cpu'})
      model.load_state_dict(checkpoint['state'])
      model.eval()
      print('model loaded.')
  
      out = model.forward(in_img, select='B')
      out = np.clip(out.detach().cpu().numpy()[0]*255, 0, 255).astype('uint8').transpose(2, 1, 0)
  
      cv2.imshow('original image', ori_img)
      cv2.imshow('network input image', img)
      cv2.imshow('result image', np.rot90(out, axes=(1, 0)))
      cv2.waitKey(0)
  
  
  if __name__ == '__main__':
      predict()
  