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) frame = logoclass.apply_frames(frame.astype(np.uint8)) 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