|
@@ -0,0 +1,118 @@
|
|
|
|
+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
|
|
|
|
+from util.norm import SpecificNorm
|
|
|
|
+import torch.nn.functional as F
|
|
|
|
+
|
|
|
|
+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,specific_person_id_nonorm,id_thres, swap_model, detect_model, save_path, temp_results_dir='./temp_results', crop_size=224, no_simswaplogo = False):
|
|
|
|
+ video_forcheck = VideoFileClip(video_path)
|
|
|
|
+ if video_forcheck.audio is None:
|
|
|
|
+ no_audio = True
|
|
|
|
+ else:
|
|
|
|
+ no_audio = False
|
|
|
|
+
|
|
|
|
+ del video_forcheck
|
|
|
|
+
|
|
|
|
+ if not no_audio:
|
|
|
|
+ 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)
|
|
|
|
+
|
|
|
|
+ spNorm =SpecificNorm()
|
|
|
|
+ mse = torch.nn.MSELoss().cuda()
|
|
|
|
+
|
|
|
|
+ # 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]
|
|
|
|
+
|
|
|
|
+ id_compare_values = []
|
|
|
|
+ frame_align_crop_tenor_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()
|
|
|
|
+
|
|
|
|
+ frame_align_crop_tenor_arcnorm = spNorm(frame_align_crop_tenor)
|
|
|
|
+ frame_align_crop_tenor_arcnorm_downsample = F.interpolate(frame_align_crop_tenor_arcnorm, scale_factor=0.5)
|
|
|
|
+ frame_align_crop_crop_id_nonorm = swap_model.netArc(frame_align_crop_tenor_arcnorm_downsample)
|
|
|
|
+
|
|
|
|
+ id_compare_values.append(mse(frame_align_crop_crop_id_nonorm,specific_person_id_nonorm).detach().cpu().numpy())
|
|
|
|
+ frame_align_crop_tenor_list.append(frame_align_crop_tenor)
|
|
|
|
+ id_compare_values_array = np.array(id_compare_values)
|
|
|
|
+ min_index = np.argmin(id_compare_values_array)
|
|
|
|
+ min_value = id_compare_values_array[min_index]
|
|
|
|
+ if min_value < id_thres:
|
|
|
|
+ swap_result = swap_model(None, frame_align_crop_tenor_list[min_index], id_vetor, None, True)[0]
|
|
|
|
+
|
|
|
|
+ reverse2wholeimage([swap_result], [frame_mat_list[min_index]], crop_size, frame, logoclass,os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)),no_simswaplogo)
|
|
|
|
+ else:
|
|
|
|
+ if not os.path.exists(temp_results_dir):
|
|
|
|
+ os.mkdir(temp_results_dir)
|
|
|
|
+ frame = frame.astype(np.uint8)
|
|
|
|
+ if not no_simswaplogo:
|
|
|
|
+ frame = logoclass.apply_frames(frame)
|
|
|
|
+ cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame)
|
|
|
|
+
|
|
|
|
+ else:
|
|
|
|
+ if not os.path.exists(temp_results_dir):
|
|
|
|
+ os.mkdir(temp_results_dir)
|
|
|
|
+ frame = frame.astype(np.uint8)
|
|
|
|
+ if not no_simswaplogo:
|
|
|
|
+ frame = logoclass.apply_frames(frame)
|
|
|
|
+ cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame)
|
|
|
|
+ else:
|
|
|
|
+ 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)
|
|
|
|
+
|
|
|
|
+ if not no_audio:
|
|
|
|
+ clips = clips.set_audio(video_audio_clip)
|
|
|
|
+
|
|
|
|
+
|
|
|
|
+ clips.write_videofile(save_path)
|
|
|
|
+
|