| ''' |
| Author: Naiyuan liu |
| Github: https://github.com/NNNNAI |
| Date: 2021-11-23 17:03:58 |
| LastEditors: Naiyuan liu |
| LastEditTime: 2021-11-24 19:19:52 |
| Description: |
| ''' |
| 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 |
| from parsing_model.model import BiSeNet |
|
|
| 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, no_simswaplogo = False,use_mask =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)) |
|
|
| |
|
|
| |
| |
| fps = video.get(cv2.CAP_PROP_FPS) |
| if os.path.exists(temp_results_dir): |
| shutil.rmtree(temp_results_dir) |
|
|
| spNorm =SpecificNorm() |
| if use_mask: |
| n_classes = 19 |
| net = BiSeNet(n_classes=n_classes) |
| net.cuda() |
| save_pth = os.path.join('./parsing_model/checkpoint', '79999_iter.pth') |
| net.load_state_dict(torch.load(save_pth)) |
| net.eval() |
| else: |
| net =None |
|
|
| |
| 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: |
| |
| 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 = [] |
| frame_align_crop_tenor_list = [] |
| for frame_align_crop in frame_align_crop_list: |
|
|
| |
| |
|
|
| 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] |
| cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame) |
| swap_result_list.append(swap_result) |
| frame_align_crop_tenor_list.append(frame_align_crop_tenor) |
|
|
| |
|
|
| reverse2wholeimage(frame_align_crop_tenor_list,swap_result_list, frame_mat_list, crop_size, frame, logoclass,\ |
| os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)),no_simswaplogo,pasring_model =net,use_mask=use_mask, norm = spNorm) |
|
|
| 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() |
|
|
| |
| 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,audio_codec='aac') |
|
|
|
|
