88 lines
2.7 KiB
Python
88 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from ultralytics import YOLO
|
|
|
|
DEFAULT_PROJECT_ROOT = Path(__file__).resolve().parents[4]
|
|
|
|
|
|
def project_root() -> Path:
|
|
raw = os.getenv("SEG_DATA_SERVER_ROOT")
|
|
if not raw:
|
|
return DEFAULT_PROJECT_ROOT
|
|
path = Path(raw).expanduser()
|
|
if path.is_absolute():
|
|
return path.resolve()
|
|
return (DEFAULT_PROJECT_ROOT / path).resolve()
|
|
|
|
|
|
PROJECT_ROOT = project_root()
|
|
|
|
MODEL_ALIASES = {
|
|
"YOLOv8n-seg": "yolov8n-seg.pt",
|
|
"YOLOv8s-seg": "yolov8s-seg.pt",
|
|
"YOLOv8m-seg": "yolov8m-seg.pt",
|
|
"YOLOv8l-seg": "yolov8l-seg.pt",
|
|
"YOLOv8x-seg": "yolov8x-seg.pt",
|
|
"YOLOv9c-seg": "yolov9c-seg.pt",
|
|
"YOLOv9e-seg": "yolov9e-seg.pt",
|
|
"YOLO11n-seg": "yolo11n-seg.pt",
|
|
"YOLO11s-seg": "yolo11s-seg.pt",
|
|
"YOLO11m-seg": "yolo11m-seg.pt",
|
|
"YOLO11l-seg": "yolo11l-seg.pt",
|
|
"YOLO11x-seg": "yolo11x-seg.pt",
|
|
}
|
|
|
|
|
|
def resolve_project_path(value: str) -> Path:
|
|
path = Path(value).expanduser()
|
|
if path.is_absolute():
|
|
return path.resolve()
|
|
return (PROJECT_ROOT / path).resolve()
|
|
|
|
|
|
def resolve_model(value: str) -> str:
|
|
candidate = MODEL_ALIASES.get(value, value)
|
|
path = Path(candidate).expanduser()
|
|
if path.is_absolute() or path.exists():
|
|
return str(path.resolve())
|
|
return candidate
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Train a YOLO segmentation model from a supplied dataset.yaml.")
|
|
parser.add_argument("--data", required=True, help="Path to YOLO dataset.yaml.")
|
|
parser.add_argument("--model", default="YOLO11n-seg", help="Model alias, weight path, or Ultralytics model name.")
|
|
parser.add_argument("--epochs", type=int, default=10)
|
|
parser.add_argument("--imgsz", type=int, default=640)
|
|
parser.add_argument("--batch", type=int, default=1)
|
|
parser.add_argument("--workers", type=int, default=0)
|
|
parser.add_argument("--device", default="cpu")
|
|
parser.add_argument("--project", required=True)
|
|
parser.add_argument("--name", default="custom_upload")
|
|
parser.add_argument("--exist-ok", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
model = YOLO(resolve_model(args.model))
|
|
result = model.train(
|
|
data=str(resolve_project_path(args.data)),
|
|
epochs=args.epochs,
|
|
imgsz=args.imgsz,
|
|
batch=args.batch,
|
|
workers=args.workers,
|
|
device=args.device,
|
|
project=str(resolve_project_path(args.project)),
|
|
name=args.name,
|
|
exist_ok=args.exist_ok,
|
|
verbose=True,
|
|
)
|
|
save_dir = getattr(result, "save_dir", None) or getattr(model.trainer, "save_dir", "")
|
|
print(f"save_dir={save_dir}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|