| 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) |
|
|
