File size: 3,397 Bytes
39732fd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import paddle
import argparse
import cv2
import numpy as np
import os
from models.model import FaceSwap, l2_norm
from models.arcface import IRBlock, ResNet
from utils.align_face import back_matrix, dealign, align_img
from utils.util import paddle2cv, cv2paddle
from utils.prepare_data import LandmarkModel
from tqdm import tqdm
def get_id_emb(id_net, id_img):
id_img = cv2.resize(id_img, (112, 112))
id_img = cv2paddle(id_img)
mean = paddle.to_tensor([[0.485, 0.456, 0.406]]).reshape((1, 3, 1, 1))
std = paddle.to_tensor([[0.229, 0.224, 0.225]]).reshape((1, 3, 1, 1))
id_img = (id_img - mean) / std
id_emb, id_feature = id_net(id_img)
id_emb = l2_norm(id_emb)
return id_emb, id_feature
def video_test(args):
paddle.set_device("gpu" if args.use_gpu else 'cpu')
faceswap_model = FaceSwap(args.use_gpu)
id_net = ResNet(block=IRBlock, layers=[3, 4, 23, 3])
id_net.set_dict(paddle.load('./checkpoints/arcface.pdparams'))
id_net.eval()
weight = paddle.load('./checkpoints/MobileFaceSwap_224.pdparams')
landmarkModel = LandmarkModel(name='landmarks')
landmarkModel.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))
id_img = cv2.imread(args.source_img_path)
#人脸检测
landmark = landmarkModel.get(id_img)
if landmark is None:
print('**** No Face Detect Error ****')
exit()
aligned_id_img, _ = align_img(id_img, landmark)
id_emb, id_feature = get_id_emb(id_net, aligned_id_img)
faceswap_model.set_model_param(id_emb, id_feature, model_weight=weight)
faceswap_model.eval()
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
cap = cv2.VideoCapture()
cap.open(args.target_video_path)
videoWriter = cv2.VideoWriter(os.path.join(args.output_path, os.path.basename(args.target_video_path)), fourcc, int(cap.get(cv2.CAP_PROP_FPS)), (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))
all_f = cap.get(cv2.CAP_PROP_FRAME_COUNT)
for i in tqdm(range(int(all_f))):
ret, frame = cap.read()
landmark = landmarkModel.get(frame)
if landmark is not None:
att_img, back_matrix = align_img(frame, landmark)
att_img = cv2paddle(att_img)
res, mask = faceswap_model(att_img)
res = paddle2cv(res)
mask = np.transpose(mask[0].numpy(), (1, 2, 0))
res = dealign(res, frame, back_matrix, mask)
frame = res
else:
print('**** No Face Detect Error ****')
videoWriter.write(frame)
cap.release()
videoWriter.release()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="MobileFaceSwap Test")
parser = argparse.ArgumentParser(description="MobileFaceSwap Test")
parser.add_argument('--source_img_path', type=str, help='path to the source image')
parser.add_argument('--target_video_path', type=str, help='path to the target video')
parser.add_argument('--output_path', type=str, default='results', help='path to the output videos')
parser.add_argument('--image_size', type=int, default=224,help='size of the test images (224 SimSwap | 256 FaceShifter)')
parser.add_argument('--merge_result', type=bool, default=True, help='output with whole image')
parser.add_argument('--use_gpu', type=bool, default=False)
args = parser.parse_args()
video_test(args) |