xunsong.li
init commit
7ccc423
raw
history blame
1.06 kB
from src.dwpose import DWposeDetector
import os
from pathlib import Path
from src.utils.util import get_fps, read_frames, save_videos_from_pil
import numpy as np
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--video_path", type=str)
args = parser.parse_args()
if not os.path.exists(args.video_path):
raise ValueError(f"Path: {args.video_path} not exists")
dir_path, video_name = (
os.path.dirname(args.video_path),
os.path.splitext(os.path.basename(args.video_path))[0],
)
out_path = os.path.join(dir_path, video_name + "_kps.mp4")
detector = DWposeDetector()
detector = detector.to(f"cuda")
fps = get_fps(args.video_path)
frames = read_frames(args.video_path)
kps_results = []
for i, frame_pil in enumerate(frames):
result, score = detector(frame_pil)
score = np.mean(score, axis=-1)
kps_results.append(result)
print(out_path)
save_videos_from_pil(kps_results, out_path, fps=fps)