64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
import time
|
|
import os
|
|
import numpy as np
|
|
import uuid
|
|
import cv2
|
|
import tqdm
|
|
import shutil
|
|
import sys
|
|
from talkingface.audio_model import AudioModel
|
|
from talkingface.render_model import RenderModel
|
|
|
|
def main():
|
|
# 检查命令行参数的数量
|
|
if len(sys.argv) != 4:
|
|
print("Usage: python demo.py <video_path> <output_video_name>")
|
|
sys.exit(1) # 参数数量不正确时退出程序
|
|
|
|
# 获取video_name参数
|
|
video_path = sys.argv[1]
|
|
print(f"Video path is set to: {video_path}")
|
|
audio_path = sys.argv[2]
|
|
print(f"Audio path is set to: {audio_path}")
|
|
output_video_name = sys.argv[3]
|
|
print(f"output video name is set to: {output_video_name}")
|
|
|
|
audioModel = AudioModel()
|
|
audioModel.loadModel("checkpoint/audio.pkl")
|
|
|
|
renderModel = RenderModel()
|
|
renderModel.loadModel("checkpoint/render.pth")
|
|
pkl_path = "{}/keypoint_rotate.pkl".format(video_path)
|
|
video_path = "{}/circle.mp4".format(video_path)
|
|
renderModel.reset_charactor(video_path, pkl_path)
|
|
|
|
# wavpath = "video_data/audio0.wav"
|
|
wavpath = audio_path
|
|
mouth_frame = audioModel.interface_wav(wavpath)
|
|
cap_input = cv2.VideoCapture(video_path)
|
|
vid_width = cap_input.get(cv2.CAP_PROP_FRAME_WIDTH) # 宽度
|
|
vid_height = cap_input.get(cv2.CAP_PROP_FRAME_HEIGHT) # 高度
|
|
cap_input.release()
|
|
task_id = str(uuid.uuid1())
|
|
os.makedirs("output/{}".format(task_id), exist_ok=True)
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|
save_path = "output/{}/silence.mp4".format(task_id)
|
|
videoWriter = cv2.VideoWriter(save_path, fourcc, 25, (int(vid_width) * 1, int(vid_height)))
|
|
for frame in tqdm.tqdm(mouth_frame):
|
|
frame = renderModel.interface(frame)
|
|
# cv2.imshow("s", frame)
|
|
# cv2.waitKey(40)
|
|
|
|
videoWriter.write(frame)
|
|
|
|
videoWriter.release()
|
|
val_video = "../output/{}.mp4".format(task_id)
|
|
os.system(
|
|
"ffmpeg -i {} -i {} -c:v libx264 -pix_fmt yuv420p -loglevel quiet {}".format(save_path, wavpath, output_video_name))
|
|
shutil.rmtree("output/{}".format(task_id))
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|