349 lines
9.9 KiB
Python
349 lines
9.9 KiB
Python
# backend/ml/train_detector_model.py
|
||
|
||
import argparse
|
||
import json
|
||
import shutil
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
|
||
from ultralytics import YOLO
|
||
from ultralytics.models.yolo.detect.train import DetectionTrainer
|
||
|
||
|
||
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 batch_sample_ids(batch):
|
||
if not isinstance(batch, dict):
|
||
return []
|
||
|
||
image_paths = batch.get("im_file")
|
||
if not image_paths:
|
||
return []
|
||
|
||
if isinstance(image_paths, (str, Path)):
|
||
image_paths = [image_paths]
|
||
|
||
out = []
|
||
seen = set()
|
||
|
||
for image_path in image_paths:
|
||
stem = Path(str(image_path)).stem.strip()
|
||
if stem and stem not in seen:
|
||
seen.add(stem)
|
||
out.append(stem)
|
||
|
||
return out
|
||
|
||
|
||
def progress_detection_trainer(train_count, val_count, train_device, fallback_epochs):
|
||
class ProgressDetectionTrainer(DetectionTrainer):
|
||
def __init__(self, *args, **kwargs):
|
||
super().__init__(*args, **kwargs)
|
||
self._preview_epoch = 0
|
||
self._preview_batch = 0
|
||
|
||
def preprocess_batch(self, batch):
|
||
epoch = int(getattr(self, "epoch", 0)) + 1
|
||
total_epochs = int(getattr(self, "epochs", fallback_epochs) or fallback_epochs)
|
||
|
||
if epoch != self._preview_epoch:
|
||
self._preview_epoch = epoch
|
||
self._preview_batch = 0
|
||
|
||
self._preview_batch += 1
|
||
|
||
# Pro Batch das gerade trainierte Bild melden – ohne Drosselung, damit
|
||
# das Frontend live immer das aktuell trainierte Bild anzeigen kann.
|
||
sample_ids = batch_sample_ids(batch)
|
||
if sample_ids:
|
||
total_batches = max(1, len(self.train_loader))
|
||
completed = (epoch - 1) + min(1.0, self._preview_batch / total_batches)
|
||
progress = 0.04 + 0.90 * (completed / max(1, total_epochs))
|
||
|
||
emit_progress(
|
||
"detector",
|
||
progress,
|
||
"Object Detector trainiert…",
|
||
epoch=epoch,
|
||
epochs=total_epochs,
|
||
sampleId=sample_ids[0],
|
||
trainSamples=train_count,
|
||
valSamples=val_count,
|
||
device=str(train_device),
|
||
)
|
||
|
||
return super().preprocess_batch(batch)
|
||
|
||
return ProgressDetectionTrainer
|
||
|
||
|
||
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
|
||
best_map50 = 0.0
|
||
best_map5095 = 0.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, best_map50, best_map5095
|
||
|
||
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
|
||
|
||
# Bestes Modell merken (YOLO speichert best.pt nach Fitness ~ mAP).
|
||
m50 = float(map50) if map50 is not None else 0.0
|
||
if m50 >= best_map50:
|
||
best_map50 = m50
|
||
if map5095 is not None:
|
||
best_map5095 = float(map5095)
|
||
|
||
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(
|
||
trainer=progress_detection_trainer(
|
||
train_count,
|
||
val_count,
|
||
train_device,
|
||
epochs,
|
||
),
|
||
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),
|
||
"trainedAt": datetime.now(timezone.utc).isoformat(),
|
||
"trainSamples": train_count,
|
||
"valSamples": val_count,
|
||
"epochs": epochs,
|
||
"imgsz": imgsz,
|
||
"device": str(train_device),
|
||
"mAP50": round(best_map50, 4),
|
||
"mAP5095": round(best_map5095, 4),
|
||
"bestEpoch": best_epoch,
|
||
}
|
||
|
||
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()
|