406 lines
11 KiB
Python
406 lines
11 KiB
Python
# backend/ml/train_pose_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.pose.train import PoseTrainer
|
|
|
|
BASE_MODEL_NAME = "yolo26n-pose.pt"
|
|
|
|
|
|
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"
|
|
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 existing_file(path):
|
|
try:
|
|
p = Path(path).expanduser().resolve()
|
|
if p.exists() and p.is_file() and p.stat().st_size > 0:
|
|
return p
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
def base_model_candidates(root):
|
|
script_dir = Path(__file__).resolve().parent
|
|
|
|
return [
|
|
script_dir / BASE_MODEL_NAME,
|
|
Path.cwd() / BASE_MODEL_NAME,
|
|
root / BASE_MODEL_NAME,
|
|
root.parent / BASE_MODEL_NAME,
|
|
root.parent.parent / BASE_MODEL_NAME,
|
|
]
|
|
|
|
|
|
def resolve_base_model(root, requested):
|
|
requested = str(requested or "").strip()
|
|
|
|
if requested and requested != BASE_MODEL_NAME:
|
|
p = existing_file(requested)
|
|
if p:
|
|
return str(p)
|
|
return requested
|
|
|
|
for candidate in base_model_candidates(root):
|
|
p = existing_file(candidate)
|
|
if p:
|
|
return str(p)
|
|
|
|
return requested or BASE_MODEL_NAME
|
|
|
|
|
|
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_pose_trainer(train_count, val_count, train_device, fallback_epochs):
|
|
class ProgressPoseTrainer(PoseTrainer):
|
|
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
|
|
|
|
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(
|
|
"pose",
|
|
progress,
|
|
"Pose 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 ProgressPoseTrainer
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--root", required=True)
|
|
parser.add_argument("--base", default="yolo26n-pose.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("--batch", default="0")
|
|
parser.add_argument("--threads", default="0")
|
|
parser.add_argument("--patience", default="20")
|
|
args = parser.parse_args()
|
|
|
|
import torch
|
|
|
|
root = Path(args.root).resolve()
|
|
dataset_root = root / "pose" / "dataset"
|
|
yaml_path = dataset_root / "dataset.yaml"
|
|
runs_dir = root / "pose" / "runs"
|
|
out_dir = root / "pose" / "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))
|
|
batch_size = max(0, safe_int(args.batch, 0))
|
|
threads = max(0, safe_int(args.threads, 0))
|
|
patience = max(0, safe_int(args.patience, 20))
|
|
|
|
if threads > 0:
|
|
torch.set_num_threads(threads)
|
|
try:
|
|
torch.set_num_interop_threads(max(1, min(threads, 4)))
|
|
except Exception:
|
|
pass
|
|
|
|
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(
|
|
"pose",
|
|
0.01,
|
|
"YOLO-Pose-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 pose train samples found")
|
|
|
|
if val_count <= 0:
|
|
raise SystemExit("no YOLO pose val samples found")
|
|
|
|
emit_progress(
|
|
"pose",
|
|
0.03,
|
|
"YOLO-Pose-Basismodell wird geladen…",
|
|
base=resolve_base_model(root, args.base),
|
|
device=str(train_device),
|
|
)
|
|
|
|
base_model = resolve_base_model(root, args.base)
|
|
model = YOLO(base_model)
|
|
|
|
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(
|
|
"pose",
|
|
0.04 + 0.90 * ((epoch - 1) / max(1, total)),
|
|
"Pose 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(
|
|
"pose",
|
|
0.04 + 0.90 * (epoch / max(1, total)),
|
|
"Pose 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 {}
|
|
|
|
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
|
|
|
|
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(
|
|
"pose",
|
|
0.04 + 0.90 * (epoch / max(1, total)),
|
|
"Pose 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(
|
|
"pose",
|
|
0.05,
|
|
"Pose Detector Training startet…",
|
|
trainSamples=train_count,
|
|
valSamples=val_count,
|
|
epochs=epochs,
|
|
imgsz=imgsz,
|
|
device=str(train_device),
|
|
)
|
|
|
|
train_kwargs = {
|
|
"trainer": progress_pose_trainer(
|
|
train_count,
|
|
val_count,
|
|
train_device,
|
|
epochs,
|
|
),
|
|
"data": str(yaml_path),
|
|
"epochs": epochs,
|
|
"imgsz": imgsz,
|
|
"project": str(runs_dir),
|
|
"name": "pose",
|
|
"exist_ok": True,
|
|
"device": train_device,
|
|
"workers": workers,
|
|
"patience": patience,
|
|
}
|
|
if batch_size > 0:
|
|
train_kwargs["batch"] = batch_size
|
|
|
|
result = model.train(**train_kwargs)
|
|
|
|
emit_progress(
|
|
"pose",
|
|
0.96,
|
|
"Bestes YOLO-Pose-Modell wird übernommen…",
|
|
lastEpoch=last_epoch,
|
|
bestEpoch=best_epoch,
|
|
device=str(train_device),
|
|
)
|
|
|
|
best = runs_dir / "pose" / "weights" / "best.pt"
|
|
last = runs_dir / "pose" / "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 / "pose"
|
|
status = {
|
|
"ok": True,
|
|
"model": str(final_model),
|
|
"sourceModel": str(best),
|
|
"baseModel": str(base_model),
|
|
"runs": str(result_path),
|
|
"trainedAt": datetime.now(timezone.utc).isoformat(),
|
|
"trainSamples": train_count,
|
|
"valSamples": val_count,
|
|
"epochs": epochs,
|
|
"imgsz": imgsz,
|
|
"batchSize": batch_size,
|
|
"patience": patience,
|
|
"threads": threads,
|
|
"workers": workers,
|
|
"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(
|
|
"pose",
|
|
1.0,
|
|
"Pose Detector fertig.",
|
|
**status,
|
|
)
|
|
|
|
print(json.dumps(status, ensure_ascii=False), flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|