File size: 2,779 Bytes
a3abb2d 1925f28 a3abb2d 1925f28 a3abb2d 1925f28 a3abb2d | 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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | from __future__ import annotations
import argparse
import shutil
from pathlib import Path
from ultralytics import YOLO
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Train a lightweight steel-surface detector and install it into the project."
)
parser.add_argument(
"--data",
default="training/data/surface_gate_detector/yolo_dataset/dataset.yaml",
help="Path to the detector dataset YAML file.",
)
parser.add_argument(
"--model",
default="yolo11n.pt",
help="Pretrained Ultralytics detector checkpoint to finetune.",
)
parser.add_argument("--epochs", type=int, default=18)
parser.add_argument("--imgsz", type=int, default=512)
parser.add_argument("--batch", type=float, default=16)
parser.add_argument("--device", default="")
parser.add_argument("--workers", type=int, default=4)
parser.add_argument("--project", default="training/runs")
parser.add_argument("--name", default="surface_gate_detector")
parser.add_argument(
"--target",
default="models/steel_surface_detector.pt",
help="Where to copy the best trained checkpoint for app inference.",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
data_path = Path(args.data).resolve()
if not data_path.exists():
raise FileNotFoundError(f"Dataset YAML not found: {data_path}")
batch = int(args.batch) if float(args.batch).is_integer() else float(args.batch)
model = YOLO(args.model)
results = model.train(
data=str(data_path),
epochs=args.epochs,
imgsz=args.imgsz,
batch=batch,
device=args.device or None,
workers=args.workers,
project=args.project,
name=args.name,
optimizer="AdamW",
patience=5,
cos_lr=True,
close_mosaic=4,
hsv_h=0.01,
hsv_s=0.15,
hsv_v=0.15,
translate=0.05,
scale=0.15,
fliplr=0.5,
mosaic=0.3,
mixup=0.0,
cache=True,
pretrained=True,
single_cls=True,
save=True,
verbose=True,
)
best_path = Path(results.save_dir) / "weights" / "best.pt"
if not best_path.exists():
raise FileNotFoundError(f"Best checkpoint was not produced at {best_path}")
target_path = Path(args.target).resolve()
target_path.parent.mkdir(parents=True, exist_ok=True)
shutil.copy2(best_path, target_path)
print(f"Training run saved to: {results.save_dir}")
print(f"Best detector copied to: {target_path}")
print("Next step: restart the FastAPI app so the steel ROI detector is used in the inspection pipeline.")
if __name__ == "__main__":
main()
|