|
import spaces |
|
import argparse |
|
|
|
import os |
|
import cv2 |
|
import torch |
|
from insightface.app import FaceAnalysis |
|
from imageio_ffmpeg import get_ffmpeg_exe |
|
|
|
@spaces.GPU |
|
def main(args): |
|
app = FaceAnalysis( |
|
providers=['CUDAExecutionProvider'], |
|
provider_options=[{'device_id': args.gpu_id}], |
|
root=args.insightface_model_path, |
|
) |
|
app.prepare(ctx_id=0, det_size=(args.height, args.width)) |
|
|
|
os.system(f'{get_ffmpeg_exe()} -i "{args.video_path}" -y -vn "{args.audio_save_path}"') |
|
|
|
kps_sequence = [] |
|
video_capture = cv2.VideoCapture(args.video_path) |
|
frame_idx = 0 |
|
while video_capture.isOpened(): |
|
ret, frame = video_capture.read() |
|
if not ret: |
|
break |
|
faces = app.get(frame) |
|
assert len(faces) == 1, f'There are {len(faces)} faces in the {frame_idx}-th frame. Only one face is supported.' |
|
|
|
kps = faces[0].kps[:3] |
|
kps_sequence.append(kps) |
|
frame_idx += 1 |
|
torch.save(kps_sequence, args.kps_sequence_save_path) |
|
|
|
|
|
if __name__ == '__main__': |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('--video_path', type=str, default='') |
|
parser.add_argument('--kps_sequence_save_path', type=str, default='') |
|
parser.add_argument('--audio_save_path', type=str, default='') |
|
parser.add_argument('--device', type=str, default='cuda') |
|
parser.add_argument('--gpu_id', type=int, default=0) |
|
parser.add_argument('--insightface_model_path', type=str, default='./model_ckpts/insightface_models/') |
|
parser.add_argument('--height', type=int, default=512) |
|
parser.add_argument('--width', type=int, default=512) |
|
args = parser.parse_args() |
|
|
|
main(args) |
|
|
