Files changed (2) hide show
  1. src/live_portrait_pipeline.py +17 -1
  2. src/utils/video.py +59 -1
src/live_portrait_pipeline.py CHANGED
@@ -10,6 +10,7 @@ torch.backends.cudnn.benchmark = True # disable CUDNN_BACKEND_EXECUTION_PLAN_DES
10
  import cv2
11
  import numpy as np
12
  import pickle
 
13
  import os.path as osp
14
  from rich.progress import track
15
 
@@ -18,7 +19,7 @@ from .config.inference_config import InferenceConfig
18
  from .config.crop_config import CropConfig
19
  from .utils.cropper import Cropper
20
  from .utils.camera import get_rotation_matrix
21
- from .utils.video import images2video, concat_frames, get_fps
22
  from .utils.crop import _transform_img, prepare_paste_back, paste_back
23
  from .utils.retargeting_utils import calc_lip_close_ratio
24
  from .utils.io import load_image_rgb, load_driving_info, resize_to_limit
@@ -177,11 +178,19 @@ class LivePortraitPipeline(object):
177
 
178
  mkdir(args.output_dir)
179
  wfp_concat = None
 
 
180
  if is_video(args.driving_info):
181
  frames_concatenated = concat_frames(I_p_lst, driving_rgb_lst, img_crop_256x256)
182
  # save (driving frames, source image, drived frames) result
183
  wfp_concat = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}_concat.mp4')
184
  images2video(frames_concatenated, wfp=wfp_concat, fps=output_fps)
 
 
 
 
 
 
185
 
186
  # save drived result
187
  wfp = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}.mp4')
@@ -190,4 +199,11 @@ class LivePortraitPipeline(object):
190
  else:
191
  images2video(I_p_lst, wfp=wfp, fps=output_fps)
192
 
 
 
 
 
 
 
 
193
  return wfp, wfp_concat
 
10
  import cv2
11
  import numpy as np
12
  import pickle
13
+ import os
14
  import os.path as osp
15
  from rich.progress import track
16
 
 
19
  from .config.crop_config import CropConfig
20
  from .utils.cropper import Cropper
21
  from .utils.camera import get_rotation_matrix
22
+ from .utils.video import images2video, concat_frames, get_fps, add_audio_to_video, has_audio_stream
23
  from .utils.crop import _transform_img, prepare_paste_back, paste_back
24
  from .utils.retargeting_utils import calc_lip_close_ratio
25
  from .utils.io import load_image_rgb, load_driving_info, resize_to_limit
 
178
 
179
  mkdir(args.output_dir)
180
  wfp_concat = None
181
+ flag_has_audio = has_audio_stream(args.driving_info)
182
+
183
  if is_video(args.driving_info):
184
  frames_concatenated = concat_frames(I_p_lst, driving_rgb_lst, img_crop_256x256)
185
  # save (driving frames, source image, drived frames) result
186
  wfp_concat = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}_concat.mp4')
187
  images2video(frames_concatenated, wfp=wfp_concat, fps=output_fps)
188
+ if flag_has_audio:
189
+ # final result with concat
190
+ wfp_concat_with_audio = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}_concat_with_audio.mp4')
191
+ add_audio_to_video(wfp_concat, args.driving_info, wfp_concat_with_audio)
192
+ os.replace(wfp_concat_with_audio, wfp_concat)
193
+ log(f"Replace {wfp_concat} with {wfp_concat_with_audio}")
194
 
195
  # save drived result
196
  wfp = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}.mp4')
 
199
  else:
200
  images2video(I_p_lst, wfp=wfp, fps=output_fps)
201
 
202
+ ######### build final result #########
203
+ if flag_has_audio:
204
+ wfp_with_audio = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}_with_audio.mp4')
205
+ add_audio_to_video(wfp, args.driving_info, wfp_with_audio)
206
+ os.replace(wfp_with_audio, wfp)
207
+ log(f"Replace {wfp} with {wfp_with_audio}")
208
+
209
  return wfp, wfp_concat
src/utils/video.py CHANGED
@@ -17,7 +17,7 @@ from .rprint import rprint as print
17
 
18
 
19
  def exec_cmd(cmd):
20
- subprocess.run(cmd, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
21
 
22
 
23
  def images2video(images, wfp, **kwargs):
@@ -143,3 +143,61 @@ def get_fps(filepath, default_fps=25):
143
  fps = default_fps
144
 
145
  return fps
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
17
 
18
 
19
  def exec_cmd(cmd):
20
+ return subprocess.run(cmd, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
21
 
22
 
23
  def images2video(images, wfp, **kwargs):
 
143
  fps = default_fps
144
 
145
  return fps
146
+
147
+
148
+ def has_audio_stream(video_path: str) -> bool:
149
+ """
150
+ Check if the video file contains an audio stream.
151
+
152
+ :param video_path: Path to the video file
153
+ :return: True if the video contains an audio stream, False otherwise
154
+ """
155
+ if osp.isdir(video_path):
156
+ return False
157
+
158
+ cmd = [
159
+ 'ffprobe',
160
+ '-v', 'error',
161
+ '-select_streams', 'a',
162
+ '-show_entries', 'stream=codec_type',
163
+ '-of', 'default=noprint_wrappers=1:nokey=1',
164
+ f'"{video_path}"'
165
+ ]
166
+
167
+ try:
168
+ # result = subprocess.run(cmd, capture_output=True, text=True)
169
+ result = exec_cmd(' '.join(cmd))
170
+ if result.returncode != 0:
171
+ log(f"Error occurred while probing video: {result.stderr}")
172
+ return False
173
+
174
+ # Check if there is any output from ffprobe command
175
+ return bool(result.stdout.strip())
176
+ except Exception as e:
177
+ log(
178
+ f"Error occurred while probing video: {video_path}, "
179
+ "you may need to install ffprobe! (https://ffmpeg.org/download.html) "
180
+ "Now set audio to false!",
181
+ style="bold red"
182
+ )
183
+ return False
184
+
185
+
186
+ def add_audio_to_video(silent_video_path: str, audio_video_path: str, output_video_path: str):
187
+ cmd = [
188
+ 'ffmpeg',
189
+ '-y',
190
+ '-i', f'"{silent_video_path}"',
191
+ '-i', f'"{audio_video_path}"',
192
+ '-map', '0:v',
193
+ '-map', '1:a',
194
+ '-c:v', 'copy',
195
+ '-shortest',
196
+ f'"{output_video_path}"'
197
+ ]
198
+
199
+ try:
200
+ exec_cmd(' '.join(cmd))
201
+ log(f"Video with audio generated successfully: {output_video_path}")
202
+ except subprocess.CalledProcessError as e:
203
+ log(f"Error occurred: {e}")