263 lines
7.0 KiB
Python
263 lines
7.0 KiB
Python
# backend/ml/train_detector_model.py
|
|
|
|
import argparse
|
|
import json
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
from ultralytics import YOLO
|
|
|
|
|
|
def emit_progress(stage, progress, message="", **extra):
|
|
out = {
|
|
"type": "progress",
|
|
"stage": stage,
|
|
"progress": max(0.0, min(1.0, float(progress))),
|
|
"message": message,
|
|
}
|
|
out.update(extra)
|
|
print(json.dumps(out, ensure_ascii=False), flush=True)
|
|
|
|
|
|
def count_yolo_samples(dataset_root: Path, split: str) -> int:
|
|
images_dir = dataset_root / "images" / split
|
|
labels_dir = dataset_root / "labels" / split
|
|
|
|
if not images_dir.exists() or not labels_dir.exists():
|
|
return 0
|
|
|
|
image_exts = {".jpg", ".jpeg", ".png", ".webp"}
|
|
count = 0
|
|
|
|
for image_path in images_dir.iterdir():
|
|
if not image_path.is_file():
|
|
continue
|
|
|
|
if image_path.suffix.lower() not in image_exts:
|
|
continue
|
|
|
|
label_path = labels_dir / f"{image_path.stem}.txt"
|
|
# Eine vorhandene, leere Labeldatei ist ein gültiges YOLO-Negativbeispiel.
|
|
if label_path.exists() and label_path.is_file():
|
|
count += 1
|
|
|
|
return count
|
|
|
|
|
|
def safe_int(value, fallback):
|
|
try:
|
|
return int(value)
|
|
except Exception:
|
|
return fallback
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--root", required=True)
|
|
parser.add_argument("--base", default="yolo26n.pt")
|
|
parser.add_argument("--epochs", default="80")
|
|
parser.add_argument("--imgsz", default="640")
|
|
parser.add_argument("--device", default="auto")
|
|
parser.add_argument("--workers", default="2")
|
|
parser.add_argument("--patience", default="20")
|
|
args = parser.parse_args()
|
|
|
|
import torch
|
|
|
|
root = Path(args.root).resolve()
|
|
dataset_root = root / "detector" / "dataset"
|
|
yaml_path = dataset_root / "dataset.yaml"
|
|
runs_dir = root / "detector" / "runs"
|
|
out_dir = root / "detector" / "model"
|
|
|
|
epochs = max(1, safe_int(args.epochs, 80))
|
|
imgsz = max(64, safe_int(args.imgsz, 640))
|
|
workers = max(0, safe_int(args.workers, 2))
|
|
patience = max(0, safe_int(args.patience, 20))
|
|
|
|
if str(args.device).lower() == "auto":
|
|
train_device = 0 if torch.cuda.is_available() else "cpu"
|
|
else:
|
|
train_device = args.device
|
|
|
|
if not yaml_path.exists():
|
|
raise SystemExit(f"dataset.yaml not found: {yaml_path}")
|
|
|
|
train_count = count_yolo_samples(dataset_root, "train")
|
|
val_count = count_yolo_samples(dataset_root, "val")
|
|
|
|
emit_progress(
|
|
"detector",
|
|
0.01,
|
|
"YOLO-Dataset wird geprüft…",
|
|
trainSamples=train_count,
|
|
valSamples=val_count,
|
|
epochs=epochs,
|
|
imgsz=imgsz,
|
|
device=str(train_device),
|
|
)
|
|
|
|
if train_count <= 0:
|
|
raise SystemExit("no YOLO train samples found")
|
|
|
|
if val_count <= 0:
|
|
raise SystemExit("no YOLO val samples found")
|
|
|
|
emit_progress(
|
|
"detector",
|
|
0.03,
|
|
"YOLO-Basismodell wird geladen…",
|
|
base=args.base,
|
|
device=str(train_device),
|
|
)
|
|
|
|
model = YOLO(args.base)
|
|
|
|
best_epoch = 0
|
|
last_epoch = 0
|
|
|
|
def on_train_epoch_start(trainer):
|
|
epoch = int(getattr(trainer, "epoch", 0)) + 1
|
|
total = int(getattr(trainer, "epochs", epochs) or epochs)
|
|
|
|
emit_progress(
|
|
"detector",
|
|
0.04 + 0.90 * ((epoch - 1) / max(1, total)),
|
|
f"Object Detector trainiert…",
|
|
epoch=epoch,
|
|
epochs=total,
|
|
trainSamples=train_count,
|
|
valSamples=val_count,
|
|
device=str(train_device),
|
|
)
|
|
|
|
def on_train_epoch_end(trainer):
|
|
nonlocal last_epoch
|
|
|
|
epoch = int(getattr(trainer, "epoch", 0)) + 1
|
|
total = int(getattr(trainer, "epochs", epochs) or epochs)
|
|
last_epoch = max(last_epoch, epoch)
|
|
|
|
emit_progress(
|
|
"detector",
|
|
0.04 + 0.90 * (epoch / max(1, total)),
|
|
f"Object Detector trainiert…",
|
|
epoch=epoch,
|
|
epochs=total,
|
|
trainSamples=train_count,
|
|
valSamples=val_count,
|
|
device=str(train_device),
|
|
)
|
|
|
|
def on_fit_epoch_end(trainer):
|
|
nonlocal best_epoch
|
|
|
|
epoch = int(getattr(trainer, "epoch", 0)) + 1
|
|
total = int(getattr(trainer, "epochs", epochs) or epochs)
|
|
metrics = getattr(trainer, "metrics", None) or {}
|
|
|
|
# Ultralytics nutzt je nach Version unterschiedliche Keys.
|
|
map50 = (
|
|
metrics.get("metrics/mAP50(B)")
|
|
or metrics.get("metrics/mAP50")
|
|
or metrics.get("mAP50")
|
|
)
|
|
map5095 = (
|
|
metrics.get("metrics/mAP50-95(B)")
|
|
or metrics.get("metrics/mAP50-95")
|
|
or metrics.get("mAP50-95")
|
|
)
|
|
|
|
if map50 is not None or map5095 is not None:
|
|
best_epoch = epoch
|
|
|
|
emit_progress(
|
|
"detector",
|
|
0.04 + 0.90 * (epoch / max(1, total)),
|
|
f"Object Detector validiert…",
|
|
epoch=epoch,
|
|
epochs=total,
|
|
mAP50=map50,
|
|
mAP5095=map5095,
|
|
device=str(train_device),
|
|
)
|
|
|
|
model.add_callback("on_train_epoch_start", on_train_epoch_start)
|
|
model.add_callback("on_train_epoch_end", on_train_epoch_end)
|
|
model.add_callback("on_fit_epoch_end", on_fit_epoch_end)
|
|
|
|
emit_progress(
|
|
"detector",
|
|
0.05,
|
|
"Object Detector Training startet…",
|
|
trainSamples=train_count,
|
|
valSamples=val_count,
|
|
epochs=epochs,
|
|
imgsz=imgsz,
|
|
device=str(train_device),
|
|
)
|
|
|
|
result = model.train(
|
|
data=str(yaml_path),
|
|
epochs=epochs,
|
|
imgsz=imgsz,
|
|
project=str(runs_dir),
|
|
name="detect",
|
|
exist_ok=True,
|
|
device=train_device,
|
|
workers=workers,
|
|
patience=patience,
|
|
)
|
|
|
|
emit_progress(
|
|
"detector",
|
|
0.96,
|
|
"Bestes YOLO-Modell wird übernommen…",
|
|
lastEpoch=last_epoch,
|
|
bestEpoch=best_epoch,
|
|
device=str(train_device),
|
|
)
|
|
|
|
best = runs_dir / "detect" / "weights" / "best.pt"
|
|
last = runs_dir / "detect" / "weights" / "last.pt"
|
|
|
|
if not best.exists():
|
|
if last.exists():
|
|
best = last
|
|
else:
|
|
raise SystemExit(f"best.pt not found after training: {best}")
|
|
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
final_model = out_dir / "best.pt"
|
|
shutil.copy2(best, final_model)
|
|
|
|
result_path = runs_dir / "detect"
|
|
status = {
|
|
"ok": True,
|
|
"model": str(final_model),
|
|
"sourceModel": str(best),
|
|
"runs": str(result_path),
|
|
"trainSamples": train_count,
|
|
"valSamples": val_count,
|
|
"epochs": epochs,
|
|
"imgsz": imgsz,
|
|
"device": str(train_device),
|
|
}
|
|
|
|
with (out_dir / "status.json").open("w", encoding="utf-8") as f:
|
|
json.dump(status, f, ensure_ascii=False, indent=2)
|
|
|
|
emit_progress(
|
|
"detector",
|
|
1.0,
|
|
"Object Detector fertig.",
|
|
**status,
|
|
)
|
|
|
|
print(json.dumps(status, ensure_ascii=False), flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|