ray96nex's picture
Upload folder using huggingface_hub
6b5b22f verified
raw
history blame contribute delete
995 Bytes
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)