import argparse
from ultralytics import YOLO

parser = argparse.ArgumentParser(description="YOLO26 岩性分类训练")
parser.add_argument("--data", default="custom_dataset", help="数据集路径")
parser.add_argument("--model", default="yolo26n-cls.pt", help="预训练权重路径")
parser.add_argument("--epochs", type=int, default=40, help="训练轮数")
parser.add_argument("--imgsz", type=int, default=224, help="输入图片尺寸")
parser.add_argument("--batch", type=int, default=32, help="批次大小")
parser.add_argument("--lr0", type=float, default=0.001, help="初始学习率")
parser.add_argument("--device", default="cuda", help="训练设备")
parser.add_argument("--patience", type=int, default=10, help="早停耐心值")
parser.add_argument("--workers", type=int, default=0, help="数据加载线程数")
parser.add_argument("--cos-lr", action="store_true", default=True, help="使用余弦学习率")
parser.add_argument("--label-smoothing", type=float, default=0.15, help="标签平滑")
parser.add_argument("--plots", action="store_true", help="绘制训练曲线")
args = parser.parse_args()

model = YOLO(args.model)

model.train(
    data=args.data,
    epochs=args.epochs,
    imgsz=args.imgsz,
    batch=args.batch,
    lr0=args.lr0,
    patience=args.patience,
    cos_lr=args.cos_lr,
    augment=True,
    label_smoothing=args.label_smoothing,
    workers=args.workers,
    task="classify",
    device=args.device,
    plots=args.plots,
    exist_ok=True,
)