53 lines
2.1 KiB
Python
53 lines
2.1 KiB
Python
"""Train a YOLO11-seg mosaic detector on a generated synthetic dataset.
|
|
|
|
Prereqs: pip install -e ".[yolo]" plus PyTorch (CUDA build recommended — see README).
|
|
|
|
Usage:
|
|
python scripts/training/train_mosaic.py --data dataset_mosaic/data.yaml --epochs 100
|
|
|
|
The resulting weights (runs/segment/<name>/weights/best.pt) drop straight into the
|
|
app: Параметры → Детектор «YOLO» → выбрать этот .pt. The class is named "mosaic",
|
|
which YoloDetector maps to CensorType.MOSAIC.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser(description="Train YOLO11-seg mosaic detector")
|
|
ap.add_argument("--data", required=True, help="path to data.yaml from the generator")
|
|
ap.add_argument("--model", default="yolo11n-seg.pt", help="base model (n/s/m...-seg)")
|
|
ap.add_argument("--epochs", type=int, default=100)
|
|
ap.add_argument("--imgsz", type=int, default=640)
|
|
ap.add_argument("--batch", default="-1", help="batch size (-1 = auto)")
|
|
ap.add_argument("--device", default=None, help="cuda / 0 / cpu (default: auto)")
|
|
ap.add_argument("--name", default="mosaic", help="run name under the project dir")
|
|
ap.add_argument("--project", default=None, help="output dir for runs (default: runs/segment)")
|
|
args = ap.parse_args()
|
|
|
|
try:
|
|
from ultralytics import YOLO
|
|
except ImportError as exc: # pragma: no cover
|
|
raise SystemExit('Не установлен ultralytics: pip install -e ".[yolo]"') from exc
|
|
|
|
batch = int(args.batch) if str(args.batch).lstrip("-").isdigit() else args.batch
|
|
model = YOLO(args.model)
|
|
results = model.train(
|
|
data=args.data,
|
|
epochs=args.epochs,
|
|
imgsz=args.imgsz,
|
|
batch=batch,
|
|
device=args.device,
|
|
name=args.name,
|
|
project=args.project,
|
|
)
|
|
save_dir = getattr(results, "save_dir", "runs/segment/" + args.name)
|
|
print(f"\nГотово. Веса: {save_dir}/weights/best.pt")
|
|
print("Подключите их в приложении: Параметры → Детектор YOLO → выбрать best.pt")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|