Some related scripts for face swaping.
Some related scripts for face swaping.
This commit is contained in:
NNNNAI
@@ -0,0 +1,55 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
# import time
|
||||
from util.add_watermark import watermark_image
|
||||
|
||||
def reverse2wholeimage(swaped_imgs, mats, crop_size, oriimg, logoclass, save_path = '',):
|
||||
|
||||
target_image_list = []
|
||||
img_mask_list = []
|
||||
for swaped_img, mat in zip(swaped_imgs, mats):
|
||||
swaped_img = swaped_img.cpu().detach().numpy().transpose((1, 2, 0))
|
||||
img_white = np.full((crop_size,crop_size), 255, dtype=float)
|
||||
|
||||
# inverse the Affine transformation matrix
|
||||
mat_rev = np.zeros([2,3])
|
||||
div1 = mat[0][0]*mat[1][1]-mat[0][1]*mat[1][0]
|
||||
mat_rev[0][0] = mat[1][1]/div1
|
||||
mat_rev[0][1] = -mat[0][1]/div1
|
||||
mat_rev[0][2] = -(mat[0][2]*mat[1][1]-mat[0][1]*mat[1][2])/div1
|
||||
div2 = mat[0][1]*mat[1][0]-mat[0][0]*mat[1][1]
|
||||
mat_rev[1][0] = mat[1][0]/div2
|
||||
mat_rev[1][1] = -mat[0][0]/div2
|
||||
mat_rev[1][2] = -(mat[0][2]*mat[1][0]-mat[0][0]*mat[1][2])/div2
|
||||
|
||||
orisize = (oriimg.shape[1], oriimg.shape[0])
|
||||
target_image = cv2.warpAffine(swaped_img, mat_rev, orisize)
|
||||
img_white = cv2.warpAffine(img_white, mat_rev, orisize)
|
||||
|
||||
|
||||
img_white[img_white>20] =255
|
||||
|
||||
img_mask = img_white
|
||||
|
||||
kernel = np.ones((10,10),np.uint8)
|
||||
img_mask = cv2.erode(img_mask,kernel,iterations = 1)
|
||||
|
||||
img_mask /= 255
|
||||
|
||||
img_mask = np.reshape(img_mask, [img_mask.shape[0],img_mask.shape[1],1])
|
||||
target_image = np.array(target_image, dtype=np.float)[..., ::-1] * 255
|
||||
|
||||
img_mask_list.append(img_mask)
|
||||
target_image_list.append(target_image)
|
||||
# target_image /= 255
|
||||
# target_image = 0
|
||||
img = np.array(oriimg, dtype=np.float)
|
||||
for img_mask, target_image in zip(img_mask_list, target_image_list):
|
||||
img = img_mask * target_image + (1-img_mask) * img
|
||||
|
||||
final_img = logoclass.apply_frames(img.astype(np.uint8))
|
||||
cv2.imwrite(save_path, final_img)
|
||||
|
||||
# cv2.imwrite('E:\\lny\\SimSwap-main\\output\\img_div.jpg', img * 255)
|
||||
# cv2.imwrite('E:\\lny\\SimSwap-main\\output\\ori_img.jpg', oriimg)
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
import os
|
||||
import cv2
|
||||
import glob
|
||||
import torch
|
||||
import shutil
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from util.reverse2original import reverse2wholeimage
|
||||
import moviepy.editor as mp
|
||||
from moviepy.editor import AudioFileClip, VideoFileClip
|
||||
from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
|
||||
import time
|
||||
from util.add_watermark import watermark_image
|
||||
|
||||
|
||||
def _totensor(array):
|
||||
tensor = torch.from_numpy(array)
|
||||
img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
|
||||
return img.float().div(255)
|
||||
|
||||
def video_swap(video_path, id_vetor, swap_model, detect_model, save_path, temp_results_dir='./temp_results', crop_size=224):
|
||||
video_audio_clip = AudioFileClip(video_path)
|
||||
video = cv2.VideoCapture(video_path)
|
||||
logoclass = watermark_image('./simswaplogo/simswaplogo.png')
|
||||
ret = True
|
||||
frame_index = 0
|
||||
|
||||
frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
|
||||
# video_WIDTH = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
|
||||
# video_HEIGHT = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
|
||||
fps = video.get(cv2.CAP_PROP_FPS)
|
||||
if os.path.exists(temp_results_dir):
|
||||
shutil.rmtree(temp_results_dir)
|
||||
|
||||
# while ret:
|
||||
for frame_index in tqdm(range(frame_count)):
|
||||
ret, frame = video.read()
|
||||
if ret:
|
||||
detect_results = detect_model.get(frame,crop_size)
|
||||
|
||||
if detect_results is not None:
|
||||
# print(frame_index)
|
||||
if not os.path.exists(temp_results_dir):
|
||||
os.mkdir(temp_results_dir)
|
||||
frame_align_crop_list = detect_results[0]
|
||||
frame_mat_list = detect_results[1]
|
||||
swap_result_list = []
|
||||
|
||||
for frame_align_crop in frame_align_crop_list:
|
||||
|
||||
# BGR TO RGB
|
||||
# frame_align_crop_RGB = frame_align_crop[...,::-1]
|
||||
|
||||
frame_align_crop_tenor = _totensor(cv2.cvtColor(frame_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
|
||||
|
||||
swap_result = swap_model(None, frame_align_crop_tenor, id_vetor, None, True)[0]
|
||||
swap_result_list.append(swap_result)
|
||||
|
||||
|
||||
|
||||
reverse2wholeimage(swap_result_list, frame_mat_list, crop_size, frame, logoclass,os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)))
|
||||
|
||||
else:
|
||||
if not os.path.exists(temp_results_dir):
|
||||
os.mkdir(temp_results_dir)
|
||||
cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame)
|
||||
else:
|
||||
break
|
||||
|
||||
# TODO,是否应该判断这个break是否是异常抛出
|
||||
video.release()
|
||||
|
||||
# image_filename_list = []
|
||||
path = os.path.join(temp_results_dir,'*.jpg')
|
||||
image_filenames = sorted(glob.glob(path))
|
||||
|
||||
clips = ImageSequenceClip(image_filenames,fps = fps)
|
||||
|
||||
final_clips = clips.set_audio(video_audio_clip)
|
||||
|
||||
# logo = (mp.ImageClip("./simswaplogo/simswap.png")
|
||||
# .set_duration(clips.duration) # 水印持续时间
|
||||
# .resize(height=100) # 水印的高度，会等比缩放
|
||||
# .margin(right=8, top=8, opacity=1) # 水印边距和透明度
|
||||
# .set_pos(("left"))) # 水印的位置
|
||||
|
||||
# final_clips = mp.CompositeVideoClip([clips, logo])
|
||||
|
||||
# final_clips.write_videofile("./output/test_beatuy_480p_full.mp4")
|
||||
final_clips.write_videofile(save_path)
|
||||
|
||||
# video = VideoFileClip(save_path)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
# video_audio_clip
|
||||
Reference in New Issue
Block a user