first commit

predict_64x64.py

@@ -0,0 +1,71 @@
+"""
+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()