File size: 995 Bytes
6b5b22f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | 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)
|