videoswap_multispecific.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133
  1. import os
  2. import cv2
  3. import glob
  4. import torch
  5. import shutil
  6. import numpy as np
  7. from tqdm import tqdm
  8. from util.reverse2original import reverse2wholeimage
  9. import moviepy.editor as mp
  10. from moviepy.editor import AudioFileClip, VideoFileClip
  11. from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
  12. import time
  13. from util.add_watermark import watermark_image
  14. from util.norm import SpecificNorm
  15. import torch.nn.functional as F
  16. def _totensor(array):
  17. tensor = torch.from_numpy(array)
  18. img = tensor.transpose(0, 1).transpose(0, 2).contiguous()
  19. return img.float().div(255)
  20. def video_swap(video_path, target_id_norm_list,source_specific_id_nonorm_list,id_thres, swap_model, detect_model, save_path, temp_results_dir='./temp_results', crop_size=224, no_simswaplogo = False):
  21. video_forcheck = VideoFileClip(video_path)
  22. if video_forcheck.audio is None:
  23. no_audio = True
  24. else:
  25. no_audio = False
  26. del video_forcheck
  27. if not no_audio:
  28. video_audio_clip = AudioFileClip(video_path)
  29. video = cv2.VideoCapture(video_path)
  30. logoclass = watermark_image('./simswaplogo/simswaplogo.png')
  31. ret = True
  32. frame_index = 0
  33. frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
  34. # video_WIDTH = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
  35. # video_HEIGHT = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
  36. fps = video.get(cv2.CAP_PROP_FPS)
  37. if os.path.exists(temp_results_dir):
  38. shutil.rmtree(temp_results_dir)
  39. spNorm =SpecificNorm()
  40. mse = torch.nn.MSELoss().cuda()
  41. # while ret:
  42. for frame_index in tqdm(range(frame_count)):
  43. ret, frame = video.read()
  44. if ret:
  45. detect_results = detect_model.get(frame,crop_size)
  46. if detect_results is not None:
  47. # print(frame_index)
  48. if not os.path.exists(temp_results_dir):
  49. os.mkdir(temp_results_dir)
  50. frame_align_crop_list = detect_results[0]
  51. frame_mat_list = detect_results[1]
  52. id_compare_values = []
  53. frame_align_crop_tenor_list = []
  54. for frame_align_crop in frame_align_crop_list:
  55. # BGR TO RGB
  56. # frame_align_crop_RGB = frame_align_crop[...,::-1]
  57. frame_align_crop_tenor = _totensor(cv2.cvtColor(frame_align_crop,cv2.COLOR_BGR2RGB))[None,...].cuda()
  58. frame_align_crop_tenor_arcnorm = spNorm(frame_align_crop_tenor)
  59. frame_align_crop_tenor_arcnorm_downsample = F.interpolate(frame_align_crop_tenor_arcnorm, scale_factor=0.5)
  60. frame_align_crop_crop_id_nonorm = swap_model.netArc(frame_align_crop_tenor_arcnorm_downsample)
  61. id_compare_values.append([])
  62. for source_specific_id_nonorm_tmp in source_specific_id_nonorm_list:
  63. id_compare_values[-1].append(mse(frame_align_crop_crop_id_nonorm,source_specific_id_nonorm_tmp).detach().cpu().numpy())
  64. frame_align_crop_tenor_list.append(frame_align_crop_tenor)
  65. id_compare_values_array = np.array(id_compare_values).transpose(1,0)
  66. min_indexs = np.argmin(id_compare_values_array,axis=0)
  67. min_value = np.min(id_compare_values_array,axis=0)
  68. swap_result_list = []
  69. swap_result_matrix_list = []
  70. for tmp_index, min_index in enumerate(min_indexs):
  71. if min_value[tmp_index] < id_thres:
  72. swap_result = swap_model(None, frame_align_crop_tenor_list[tmp_index], target_id_norm_list[min_index], None, True)[0]
  73. swap_result_list.append(swap_result)
  74. swap_result_matrix_list.append(frame_mat_list[tmp_index])
  75. else:
  76. pass
  77. if len(swap_result_list) !=0:
  78. reverse2wholeimage(swap_result_list, swap_result_matrix_list, crop_size, frame, logoclass,os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)),no_simswaplogo)
  79. else:
  80. if not os.path.exists(temp_results_dir):
  81. os.mkdir(temp_results_dir)
  82. frame = frame.astype(np.uint8)
  83. if not no_simswaplogo:
  84. frame = logoclass.apply_frames(frame)
  85. cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame)
  86. else:
  87. if not os.path.exists(temp_results_dir):
  88. os.mkdir(temp_results_dir)
  89. frame = frame.astype(np.uint8)
  90. if not no_simswaplogo:
  91. frame = logoclass.apply_frames(frame)
  92. cv2.imwrite(os.path.join(temp_results_dir, 'frame_{:0>7d}.jpg'.format(frame_index)), frame)
  93. else:
  94. break
  95. video.release()
  96. # image_filename_list = []
  97. path = os.path.join(temp_results_dir,'*.jpg')
  98. image_filenames = sorted(glob.glob(path))
  99. clips = ImageSequenceClip(image_filenames,fps = fps)
  100. if not no_audio:
  101. clips = clips.set_audio(video_audio_clip)
  102. clips.write_videofile(save_path)