import argparse
from ultralytics import YOLO

def train_model(data_path, epochs=100, imgsz=640, batch=16, model_type='yolo11n-pose.pt', name='sprint_pose'):
    """Train a YOLO pose model."""
    model = YOLO(model_type)
    
    results = model.train(
        data=data_path,
        epochs=epochs,
        imgsz=imgsz,
        batch=batch,
        name=name
    )
    print(f"✅ Training complete. Results saved in runs/pose/{name}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", required=True, help="Path to data.yaml")
    parser.add_argument("--epochs", type=int, default=100)
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--batch", type=int, default=16)
    parser.add_argument("--model", default='yolo11n-pose.pt')
    parser.add_argument("--name", default='sprint_pose')
    args = parser.parse_args()
    
    train_model(args.data, args.epochs, args.imgsz, args.batch, args.model, args.name)