move extract kps_sequence from video as util

app.py

@@ -3,6 +3,7 @@ import shutil
 import subprocess
 
 from inference import InferenceEngine
+from sequence_utils import extract_kps_sequence_from_video
 
 output_dir = "output"
 temp_audio_path = "temp.mp3"
@@ -73,13 +74,13 @@ def run_demo(
         progress((25,100), desc="Extract keypoints and audio...")
         audio_path = video.replace(".mp4", ".mp3")
 
-        subprocess.run([
-            "python", 
-            "scripts/extract_kps_sequence_and_audio.py",
-            "--video_path", video,
-            "--kps_sequence_save_path", kps_sequence_save_path,
-            "--audio_save_path", audio_path
-        ], check=True)
+        extract_kps_sequence_from_video(
+            INFERENCE_ENGINE.app,
+            video,
+            audio_path,
+            kps_sequence_save_path
+        )
+
         progress((50,100), desc="Keypoints and audio extracted successfully.")
         #return "Keypoints and audio extracted successfully."
         rem_progress = (75,100)

scripts/extract_kps_sequence_and_audio.py

@@ -1,47 +0,0 @@
-import argparse
-
-import os
-import cv2
-import torch
-from insightface.app import FaceAnalysis
-from imageio_ffmpeg import get_ffmpeg_exe
-
-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)

sequence_utils.py

@@ -0,0 +1,26 @@
+import os
+import cv2
+import torch
+from imageio_ffmpeg import get_ffmpeg_exe
+
+def extract_kps_sequence_from_video(
+        face_analysis_app,
+        video_path, audio_save_path,
+        kps_sequence_save_path):
+    os.system(f'{get_ffmpeg_exe()} -i "{video_path}" -y -vn "{audio_save_path}"')
+
+    kps_sequence = []
+    video_capture = cv2.VideoCapture(video_path)
+    frame_idx = 0
+    while video_capture.isOpened():
+        ret, frame = video_capture.read()
+        if not ret:
+            break
+        faces = face_analysis_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, kps_sequence_save_path)
+