Some related scripts for face swaping.

Some related scripts for face swaping.

util/reverse2original.py

@@ -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)
+

util/videoswap.py

@@ -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