example/run_full_pipeline.py

import os
import sys
import argparse
from video_utils import extract_frames, create_video_from_frames

def run_pipeline(video_path, model_path, output_dir, sigma=2.0):
    """Orchestrate the full pose estimation pipeline."""
    # 1. Setup paths
    frames_dir = os.path.join(output_dir, "frames")
    labels_dir = os.path.join(output_dir, "labels_raw")
    annotated_dir = os.path.join(output_dir, "annotated_frames")
    keypoints_csv = os.path.join(output_dir, "yolo_keypoints.csv")
    final_video = os.path.join(output_dir, "annotated_output.mp4")
    
    os.makedirs(output_dir, exist_ok=True)

    # 2. Extract frames
    print("--- Step 1: Extracting Frames ---")
    fps, width, height, num_frames = extract_frames(video_path, frames_dir)

    # 3. Inference
    print("\n--- Step 2: Running YOLO Inference ---")
    os.system(f'{sys.executable} inference_yolo_pose.py --frames "{frames_dir}" --labels "{labels_dir}" --model "{model_path}" --width {width} --height {height}')

    # 4. Post-processing
    print("\n--- Step 3: Post-processing Keypoints ---")
    os.system(f'{sys.executable} post_process_keypoints.py --labels "{labels_dir}" --output "{keypoints_csv}" --width {width} --height {height} --sigma {sigma}')

    # 5. Annotation
    print("\n--- Step 4: Annotating Frames ---")
    os.system(f'{sys.executable} annotate_video.py --frames "{frames_dir}" --csv "{keypoints_csv}" --output "{annotated_dir}"')

    # 6. Final Video
    print("\n--- Step 5: Creating Final Video ---")
    create_video_from_frames(annotated_dir, final_video, fps, width, height)
    
    print(f"\n✅ Pipeline complete! Output saved to {output_dir}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--video", required=True, help="Input video path")
    parser.add_argument("--model", required=True, help="YOLO model weights path")
    parser.add_argument("--output", default="video_output", help="Output directory")
    parser.add_argument("--sigma", type=float, default=2.0, help="Smoothing sigma")
    args = parser.parse_args()
    
    run_pipeline(args.video, args.model, args.output, args.sigma)