tools/vid2pose.py

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)