Generate output video directly to disk

This commit is contained in:
AleD 2024-07-13 20:50:32 +09:00
parent 4c643a86b6
commit 41b4430b5e
2 changed files with 12 additions and 8 deletions

View File

@ -220,7 +220,7 @@ class LivePortraitPipeline(object):
# driving frame | source image | generation, or source image | generation
frames_concatenated = concat_frames(driving_rgb_crop_256x256_lst, img_crop_256x256, I_p_lst)
wfp_concat = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}_concat.mp4')
images2video(frames_concatenated, wfp=wfp_concat, fps=output_fps)
images2video(frames_concatenated, wfp=wfp_concat, fps=output_fps, frames_count=len(driving_rgb_crop_256x256_lst))
if flag_has_audio:
# final result with concact

View File

@ -37,12 +37,18 @@ def images2video(images, wfp, **kwargs):
codec=codec, quality=quality, ffmpeg_params=ffmpeg_params, pixelformat=pixelformat, macro_block_size=macro_block_size
)
n = len(images)
n = len(images) if hasattr(images, '__len__') else kwargs.get('frames_count')
img_it = iter(images)
for i in track(range(n), description='Writing', transient=True):
try:
img = next(img_it)
except StopIteration:
break
if image_mode.lower() == 'bgr':
writer.append_data(images[i][..., ::-1])
writer.append_data(img[..., ::-1])
else:
writer.append_data(images[i])
writer.append_data(img)
writer.close()
@ -82,10 +88,9 @@ def blend(img: np.ndarray, mask: np.ndarray, background_color=(255, 255, 255)):
def concat_frames(driving_image_lst, source_image, I_p_lst):
# TODO: add more concat style, e.g., left-down corner driving
out_lst = []
h, w, _ = I_p_lst[0].shape
for idx, _ in track(enumerate(I_p_lst), total=len(I_p_lst), description='Concatenating result...'):
for idx, _ in enumerate(I_p_lst):
I_p = I_p_lst[idx]
source_image_resized = cv2.resize(source_image, (w, h))
@ -96,8 +101,7 @@ def concat_frames(driving_image_lst, source_image, I_p_lst):
driving_image_resized = cv2.resize(driving_image, (w, h))
out = np.hstack((driving_image_resized, source_image_resized, I_p))
out_lst.append(out)
return out_lst
yield out
class VideoWriter: