56 lines
1.4 KiB
Python
56 lines
1.4 KiB
Python
# backend\ml\train_detector_model.py
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from ultralytics import YOLO
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--root", required=True)
|
|
parser.add_argument("--base", default="yolo11n.pt")
|
|
parser.add_argument("--epochs", default="80")
|
|
parser.add_argument("--imgsz", default="640")
|
|
args = parser.parse_args()
|
|
|
|
root = Path(args.root).resolve()
|
|
yaml_path = root / "detector" / "dataset" / "dataset.yaml"
|
|
runs_dir = root / "detector" / "runs"
|
|
|
|
if not yaml_path.exists():
|
|
raise SystemExit(f"dataset.yaml not found: {yaml_path}")
|
|
|
|
model = YOLO(args.base)
|
|
|
|
result = model.train(
|
|
data=str(yaml_path),
|
|
epochs=int(args.epochs),
|
|
imgsz=int(args.imgsz),
|
|
project=str(runs_dir),
|
|
name="detect",
|
|
exist_ok=True,
|
|
device="cpu",
|
|
workers=2,
|
|
patience=20,
|
|
)
|
|
|
|
best = runs_dir / "detect" / "weights" / "best.pt"
|
|
if not best.exists():
|
|
raise SystemExit(f"best.pt not found after training: {best}")
|
|
|
|
out_dir = root / "detector" / "model"
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
final_model = out_dir / "best.pt"
|
|
final_model.write_bytes(best.read_bytes())
|
|
|
|
print(json.dumps({
|
|
"ok": True,
|
|
"model": str(final_model),
|
|
"runs": str(runs_dir / "detect"),
|
|
}))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |