From 78d560d3cadbfde32b6a1f1632d1c3d023ba4efc Mon Sep 17 00:00:00 2001 From: loootte <46289941+loootte@users.noreply.github.com> Date: Sat, 1 Aug 2026 09:55:21 +0800 Subject: [PATCH] feat(train): Streamlit UI for layout import, train, and eval (#101) Add ui/app.py and ui_backend for selecting .enpu.json, importing layout GT, launching train/eval subprocesses, and showing metrics/logs/history. Closes #101 --- README.md | 2 +- train/README.md | 39 ++- train/enpu_train/ui_backend.py | 542 +++++++++++++++++++++++++++++++++ train/requirements.txt | 4 +- train/scripts/run_ui.py | 28 ++ train/tests/test_ui_backend.py | 83 +++++ train/ui/app.py | 419 +++++++++++++++++++++++++ 7 files changed, 1108 insertions(+), 9 deletions(-) create mode 100644 train/enpu_train/ui_backend.py create mode 100644 train/scripts/run_ui.py create mode 100644 train/tests/test_ui_backend.py create mode 100644 train/ui/app.py diff --git a/README.md b/README.md index b1c1216..4bca5d7 100644 --- a/README.md +++ b/README.md @@ -62,7 +62,7 @@ L5 音符节点 音高数字 OCR（+几何兜底）+ 时值线 / 高低音点 桌面在结构模式下可叠图查看 L1–L5，L3 以分割线编辑为主。 完整说明：[architecture-structure-first.md](./docs/architecture-structure-first.md) · [l3-split-model.md](./docs/l3-split-model.md) · [architecture.md](./docs/architecture.md)。 -L1–L3 **布局训练**（#92–#95）：数据规范 [docs/train/l1-l3-data-spec.md](./docs/train/l1-l3-data-spec.md) · 模型方案 [l1-l3-model-design.md](./docs/train/l1-l3-model-design.md) · Framework [`train/`](./train/)。 +L1–L3 **布局训练**（#92–#95 / UI #101）：数据规范 [docs/train/l1-l3-data-spec.md](./docs/train/l1-l3-data-spec.md) · 模型方案 [l1-l3-model-design.md](./docs/train/l1-l3-model-design.md) · Framework + UI [`train/`](./train/)（`python scripts/run_ui.py`）。 --- diff --git a/train/README.md b/train/README.md index 13592fa..6926d80 100644 --- a/train/README.md +++ b/train/README.md @@ -16,18 +16,17 @@ train/ requirements.txt configs/mvp_l2_l3.yaml enpu_train/ - data/ # Dataset + 合成样本 - models/ # L2 page y-heat + L3 row x-heat - losses/ - metrics/ # L2 IoU + L3 split count / mean_abs_x - engine/ # train / eval - export/ # state_dict + ONNX + data/ models/ losses/ metrics/ engine/ export/ + ui_backend.py # #101 UI 调用层 viz.py + ui/ + app.py # Streamlit 训练 UI (#101) scripts/ train.py eval.py viz_sample.py export_from_enpu_project.py + run_ui.py tests/ ``` @@ -42,6 +41,30 @@ pip install -r requirements.txt 需要 **Python 3.10+**、**PyTorch**。有 GPU 时可在配置里设 `train.device: cuda`。 +## 训练 UI（#101） + +图形界面：选择 `.enpu.json` → 导入 Layout → 一键训练 → 一键测试 → 看指标/日志。 + +```powershell +cd train +.\.venv\Scripts\Activate.ps1 +# 需能 import 仓库 core（layout 导出） +$env:PYTHONPATH = "$PWD;$PWD\..\core" +python scripts\run_ui.py +# 或: streamlit run ui/app.py +``` + +浏览器打开终端提示的本地 URL（默认 `http://localhost:8501`）。 + +| 面板 | 功能 | +|------|------| +| 数据集 | 填写工程路径，导入到 `samples/layout/`；列表校验状态 | +| 训练 | 任务/epochs/device，子进程训练，刷新进度与 loss 曲线 | +| 测试 | 选 ckpt + 数据目录，展示 L2 IoU / L3 x 误差 | +| 历史 | 扫描 `runs/` | + +**注意：** 改框/拖线仍在恩谱**桌面**完成；本 UI 只负责导入与训练闭环。 + ## 数据准备 1. **真实工程** → layout 样本（#93）： @@ -109,6 +132,7 @@ python scripts\eval.py --ckpt runs\mvp_l2_l3\best.pt --data ..\samples\layout ```powershell cd train +$env:PYTHONPATH = "$PWD;$PWD\..\core" python -m pytest tests -q ``` @@ -117,5 +141,6 @@ python -m pytest tests -q - 大规模真实集 / 完整合成流水线 - core 内完整 `learned_l1l3` 推理插件 - 精度超过几何基线的承诺 +- 训练 UI 内嵌完整标注编辑器（#101 非目标） -父任务：[Issue #92](https://github.com/loootte/EnPu/issues/92) · 本任务 [#95](https://github.com/loootte/EnPu/issues/95) +相关：父任务 [#92](https://github.com/loootte/EnPu/issues/92) · Framework [#95](https://github.com/loootte/EnPu/issues/95) · UI [#101](https://github.com/loootte/EnPu/issues/101) diff --git a/train/enpu_train/ui_backend.py b/train/enpu_train/ui_backend.py new file mode 100644 index 0000000..52e217d --- /dev/null +++ b/train/enpu_train/ui_backend.py @@ -0,0 +1,542 @@ +"""Backend helpers for the training UI (#101). + +Thin wrappers around layout_gt export + train/eval scripts. +UI must not embed training algorithms — only call these entry points. +""" + +from __future__ import annotations + +import json +import os +import signal +import subprocess +import sys +import time +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +TRAIN_ROOT = Path(__file__).resolve().parents[1] +REPO_ROOT = TRAIN_ROOT.parent +CORE_ROOT = REPO_ROOT / "core" +DEFAULT_LAYOUT_ROOT = REPO_ROOT / "samples" / "layout" +DEFAULT_RUNS = TRAIN_ROOT / "runs" +JOBS_DIR = DEFAULT_RUNS / "ui_jobs" + + +def _ensure_sys_path() -> None: + for p in (str(TRAIN_ROOT), str(CORE_ROOT)): + if p not in sys.path: + sys.path.insert(0, p) + + +def discover_layout_samples(roots: list[str | Path] | None = None) -> list[Path]: + roots = roots or [DEFAULT_LAYOUT_ROOT, TRAIN_ROOT / "data_cache" / "synth"] + found: list[Path] = [] + for r in roots: + p = Path(r) + if not p.is_dir(): + continue + for layout in p.rglob("layout.json"): + found.append(layout.parent) + # unique, stable order + uniq = sorted({f.resolve() for f in found}, key=lambda x: str(x).lower()) + return uniq + + +@dataclass +class SampleInfo: + path: str + sample_id: str + ok: bool + errors: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + n_systems: int = 0 + n_splits: int = 0 + n_measures: int = 0 + image_path: str | None = None + width: int | None = None + height: int | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def inspect_layout_sample(sample_dir: str | Path) -> SampleInfo: + _ensure_sys_path() + from app.layout_gt.validate import validate_layout_sample + + sample_dir = Path(sample_dir) + layout_path = sample_dir / "layout.json" + info = SampleInfo( + path=str(sample_dir), + sample_id=sample_dir.name, + ok=False, + ) + if not layout_path.is_file(): + info.errors = ["missing layout.json"] + return info + try: + data = json.loads(layout_path.read_text(encoding="utf-8")) + except Exception as e: + info.errors = [f"invalid JSON: {e}"] + return info + + info.sample_id = str(data.get("id") or sample_dir.name) + img = data.get("image") or {} + info.width = img.get("width") + info.height = img.get("height") + rel = img.get("path") or "image.png" + ip = sample_dir / rel + if ip.is_file(): + info.image_path = str(ip) + else: + for cand in sample_dir.glob("image.*"): + info.image_path = str(cand) + break + + systems = (data.get("l2") or {}).get("systems") or [] + rows = (data.get("l3") or {}).get("rows") or [] + info.n_systems = len(systems) + info.n_splits = sum(len(r.get("splits") or []) for r in rows) + info.n_measures = sum(len(r.get("measures") or []) for r in rows) + + result = validate_layout_sample(data) + info.ok = result.ok + info.errors = list(result.errors) + info.warnings = list(result.warnings) + return info + + +def list_samples_info(roots: list[str | Path] | None = None) -> list[SampleInfo]: + return [inspect_layout_sample(p) for p in discover_layout_samples(roots)] + + +def import_enpu_project( + project_path: str | Path, + *, + out_dir: str | Path | None = None, + sample_id: str | None = None, +) -> dict[str, Any]: + """Export .enpu.json → layout sample dir. Returns sample info + paths.""" + _ensure_sys_path() + from app.layout_gt.export import export_project_to_sample_dir + + project_path = Path(project_path) + if not project_path.is_file(): + raise FileNotFoundError(f"project not found: {project_path}") + + sid = sample_id or project_path.stem.replace(".enpu", "") + # sanitize dir name + safe = "".join(c if c.isalnum() or c in "-_" else "_" for c in sid)[:80] + if not safe: + safe = f"import_{int(time.time())}" + out = Path(out_dir) if out_dir else (DEFAULT_LAYOUT_ROOT / safe) + out.mkdir(parents=True, exist_ok=True) + + try: + sample = export_project_to_sample_dir( + project_path, + out, + sample_id=sid, + copy_image=True, + validate=True, + ) + err = None + except ValueError as e: + # still write partial if possible — re-raise with message + raise ValueError(str(e)) from e + + info = inspect_layout_sample(out) + return { + "ok": info.ok, + "out_dir": str(out), + "sample_id": sample.get("id") if isinstance(sample, dict) else sid, + "info": info.to_dict(), + "error": err, + } + + +def list_runs(runs_root: str | Path | None = None) -> list[dict[str, Any]]: + root = Path(runs_root or DEFAULT_RUNS) + if not root.is_dir(): + return [] + runs: list[dict[str, Any]] = [] + for d in sorted(root.iterdir(), key=lambda p: p.stat().st_mtime, reverse=True): + if not d.is_dir() or d.name == "ui_jobs": + continue + meta = { + "name": d.name, + "path": str(d), + "mtime": datetime.fromtimestamp(d.stat().st_mtime, tz=timezone.utc).isoformat(), + "has_best": (d / "best.pt").is_file(), + "has_last": (d / "last.pt").is_file(), + "has_history": (d / "history.json").is_file(), + } + hist = d / "history.json" + if hist.is_file(): + try: + h = json.loads(hist.read_text(encoding="utf-8")) + meta["epochs"] = len(h) if isinstance(h, list) else None + if isinstance(h, list) and h: + last = h[-1] + meta["last_train_loss"] = last.get("train_loss") + val = last.get("val") or {} + meta["last_val"] = val + except Exception: + pass + job = d / "job.json" + if job.is_file(): + try: + meta["job"] = json.loads(job.read_text(encoding="utf-8")) + except Exception: + pass + runs.append(meta) + return runs + + +def load_history(run_dir: str | Path) -> list[dict[str, Any]]: + p = Path(run_dir) / "history.json" + if not p.is_file(): + return [] + data = json.loads(p.read_text(encoding="utf-8")) + return data if isinstance(data, list) else [] + + +def load_eval_metrics(path: str | Path) -> dict[str, Any] | None: + p = Path(path) + if not p.is_file(): + return None + return json.loads(p.read_text(encoding="utf-8")) + + +@dataclass +class TrainJobSpec: + run_name: str + tasks: list[str] = field(default_factory=lambda: ["l2", "l3"]) + epochs: int = 2 + batch_size: int = 2 + lr: float = 1e-3 + device: str = "cpu" + data_roots: list[str] = field(default_factory=list) + synth_count: int = 4 + val_ratio: float = 0.25 + skip_export: bool = False + + +def write_train_config(spec: TrainJobSpec, out_dir: Path) -> Path: + """Write a yaml config for scripts/train.py.""" + out_dir.mkdir(parents=True, exist_ok=True) + roots = list(spec.data_roots) if spec.data_roots else [str(DEFAULT_LAYOUT_ROOT)] + # use forward-friendly paths as strings + cfg = { + "tasks": list(spec.tasks), + "data": { + "roots": roots, + "synth_count": int(spec.synth_count), + "synth_dir": str(TRAIN_ROOT / "data_cache" / "synth"), + "page_size": [384, 512], + "row_size": [64, 256], + "l2_heat_len": 128, + "l3_heat_len": 128, + "augment": True, + "val_ratio": float(spec.val_ratio), + }, + "train": { + "epochs": int(spec.epochs), + "batch_size": int(spec.batch_size), + "lr": float(spec.lr), + "weight_decay": 0.0001, + "l2_loss_weight": 1.0, + "l3_loss_weight": 1.5, + "device": spec.device, + "num_workers": 0, + "out_dir": str(out_dir), + "log_every": 1, + }, + "export": { + "state_dict": str(out_dir / "export" / "layout_net.pt"), + "onnx_dir": str(out_dir / "export" / "onnx"), + }, + } + try: + import yaml + + text = yaml.safe_dump(cfg, allow_unicode=True, sort_keys=False) + except Exception: + text = json.dumps(cfg, ensure_ascii=False, indent=2) + cfg_path = out_dir / "config.json" + cfg_path.write_text(text, encoding="utf-8") + return cfg_path + + cfg_path = out_dir / "config.yaml" + cfg_path.write_text(text, encoding="utf-8") + return cfg_path + + +def start_train_job(spec: TrainJobSpec) -> dict[str, Any]: + """Launch train.py as a subprocess; returns job metadata (Windows-safe logging).""" + JOBS_DIR.mkdir(parents=True, exist_ok=True) + run_dir = DEFAULT_RUNS / spec.run_name + run_dir.mkdir(parents=True, exist_ok=True) + cfg_path = write_train_config(spec, run_dir) + log_path = run_dir / "train.log" + job_path = run_dir / "job.json" + + cmd = [ + sys.executable, + "-u", + str(TRAIN_ROOT / "scripts" / "train.py"), + "--config", + str(cfg_path), + "--device", + spec.device, + "--epochs", + str(spec.epochs), + ] + if spec.skip_export: + cmd.append("--skip-export") + + env = os.environ.copy() + pp = [str(TRAIN_ROOT), str(CORE_ROOT)] + if env.get("PYTHONPATH"): + pp.append(env["PYTHONPATH"]) + env["PYTHONPATH"] = os.pathsep.join(pp) + env["PYTHONUNBUFFERED"] = "1" + + # Keep log file handle open for the child process lifetime (do not close here). + log_f = open(log_path, "w", encoding="utf-8", errors="replace") + kwargs: dict[str, Any] = { + "cwd": str(TRAIN_ROOT), + "stdout": log_f, + "stderr": subprocess.STDOUT, + "env": env, + } + if os.name == "nt": + kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP + else: + kwargs["start_new_session"] = True + proc = subprocess.Popen(cmd, **kwargs) + + job = { + "kind": "train", + "pid": proc.pid, + "cmd": cmd, + "run_dir": str(run_dir), + "log_path": str(log_path), + "cfg_path": str(cfg_path), + "started_at": datetime.now(timezone.utc).isoformat(), + "status": "running", + "spec": asdict(spec), + } + job_path.write_text(json.dumps(job, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + return job + + +def _pid_running(pid: int) -> bool: + if pid <= 0: + return False + if os.name == "nt": + try: + import ctypes + + kernel32 = ctypes.windll.kernel32 # type: ignore[attr-defined] + PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 + handle = kernel32.OpenProcess( + PROCESS_QUERY_LIMITED_INFORMATION, False, int(pid) + ) + if handle: + kernel32.CloseHandle(handle) + return True + return False + except Exception: + try: + out = subprocess.run( + ["tasklist", "/FI", f"PID eq {pid}"], + capture_output=True, + timeout=10, + ) + text = out.stdout.decode("gbk", errors="replace") + return str(pid) in text + except Exception: + return False + try: + os.kill(pid, 0) + return True + except OSError: + return False + + +def poll_train_job(run_dir: str | Path) -> dict[str, Any]: + run_dir = Path(run_dir) + job_path = run_dir / "job.json" + status: dict[str, Any] = { + "run_dir": str(run_dir), + "status": "unknown", + "history": [], + "log_tail": "", + "pid": None, + } + if job_path.is_file(): + try: + job = json.loads(job_path.read_text(encoding="utf-8")) + status["pid"] = job.get("pid") + status["job"] = job + except Exception as e: + status["error"] = f"job.json: {e}" + return status + else: + status["status"] = "no_job" + return status + + pid = status.get("pid") + running = _pid_running(int(pid)) if pid else False + history = load_history(run_dir) + status["history"] = history + log_path = run_dir / "train.log" + if log_path.is_file(): + try: + text = log_path.read_text(encoding="utf-8", errors="replace") + status["log_tail"] = "\n".join(text.splitlines()[-80:]) + status["log_path"] = str(log_path) + except Exception: + pass + + if running: + status["status"] = "running" + else: + # finished — success if history and best/last exist + if (run_dir / "best.pt").is_file() or (run_dir / "last.pt").is_file(): + status["status"] = "succeeded" + elif history: + status["status"] = "succeeded" + else: + # Distinguish "still starting" (empty log, just spawned) vs failed + log_path = run_dir / "train.log" + age = time.time() - run_dir.stat().st_mtime + log_size = log_path.stat().st_size if log_path.is_file() else 0 + started = (status.get("job") or {}).get("started_at") + if log_size == 0 and age < 15: + status["status"] = "starting" + else: + status["status"] = "failed" + # update job.json only on terminal states + if status["status"] in ("succeeded", "failed", "cancelled"): + try: + job = status.get("job") or {} + job["status"] = status["status"] + job["finished_at"] = datetime.now(timezone.utc).isoformat() + job_path.write_text( + json.dumps(job, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + except Exception: + pass + return status + + +def cancel_train_job(run_dir: str | Path) -> dict[str, Any]: + run_dir = Path(run_dir) + job_path = run_dir / "job.json" + if not job_path.is_file(): + return {"ok": False, "error": "no job.json"} + job = json.loads(job_path.read_text(encoding="utf-8")) + pid = int(job.get("pid") or 0) + if not pid: + return {"ok": False, "error": "no pid"} + try: + if os.name == "nt": + subprocess.run( + ["taskkill", "/PID", str(pid), "/T", "/F"], + capture_output=True, + text=True, + timeout=30, + ) + else: + os.killpg(pid, signal.SIGTERM) + except Exception as e: + return {"ok": False, "error": str(e), "pid": pid} + job["status"] = "cancelled" + job["finished_at"] = datetime.now(timezone.utc).isoformat() + job_path.write_text(json.dumps(job, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + return {"ok": True, "pid": pid} + + +def run_eval( + ckpt: str | Path, + data_root: str | Path, + *, + out_json: str | Path | None = None, + device: str = "cpu", +) -> dict[str, Any]: + """Synchronous eval via scripts/eval.py (usually fast on toy data).""" + ckpt = Path(ckpt) + data_root = Path(data_root) + if not ckpt.is_file(): + raise FileNotFoundError(f"ckpt not found: {ckpt}") + if out_json is None: + out_json = ckpt.parent / "eval_ui.json" + out_json = Path(out_json) + + cmd = [ + sys.executable, + "-u", + str(TRAIN_ROOT / "scripts" / "eval.py"), + "--ckpt", + str(ckpt), + "--data", + str(data_root), + "--device", + device, + "--out", + str(out_json), + ] + env = os.environ.copy() + pp = [str(TRAIN_ROOT), str(CORE_ROOT)] + if env.get("PYTHONPATH"): + pp.append(env["PYTHONPATH"]) + env["PYTHONPATH"] = os.pathsep.join(pp) + + proc = subprocess.run( + cmd, + cwd=str(TRAIN_ROOT), + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + env=env, + timeout=600, + ) + result: dict[str, Any] = { + "returncode": proc.returncode, + "stdout": proc.stdout, + "stderr": proc.stderr, + "out_json": str(out_json), + } + if out_json.is_file(): + result["metrics"] = json.loads(out_json.read_text(encoding="utf-8")) + if proc.returncode != 0: + result["ok"] = False + result["error"] = (proc.stderr or proc.stdout or f"exit {proc.returncode}")[-2000:] + else: + result["ok"] = True + return result + + +def default_ckpt_for_run(run_dir: str | Path) -> Path | None: + run_dir = Path(run_dir) + for name in ("best.pt", "last.pt"): + p = run_dir / name + if p.is_file(): + return p + return None + + +def cuda_available() -> bool: + try: + import torch + + return bool(torch.cuda.is_available()) + except Exception: + return False diff --git a/train/requirements.txt b/train/requirements.txt index a52f2ed..40b7a83 100644 --- a/train/requirements.txt +++ b/train/requirements.txt @@ -1,7 +1,9 @@ -# EnPu train framework (#95) +# EnPu train framework (#95) + UI (#101) # Install in a venv separate from core if preferred. +# ASCII-only file: Windows pip may decode requirements as GBK. torch>=2.0 numpy>=1.24 Pillow>=10.0 PyYAML>=6.0 tqdm>=4.65 +streamlit>=1.28 diff --git a/train/scripts/run_ui.py b/train/scripts/run_ui.py new file mode 100644 index 0000000..fbaa1ca --- /dev/null +++ b/train/scripts/run_ui.py @@ -0,0 +1,28 @@ +#!/usr/bin/env python3 +"""Launch Streamlit training UI (#101).""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +TRAIN_ROOT = Path(__file__).resolve().parents[1] +APP = TRAIN_ROOT / "ui" / "app.py" + + +def main() -> int: + cmd = [ + sys.executable, + "-m", + "streamlit", + "run", + str(APP), + "--server.headless", + "true", + ] + return subprocess.call(cmd, cwd=str(TRAIN_ROOT)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/train/tests/test_ui_backend.py b/train/tests/test_ui_backend.py new file mode 100644 index 0000000..b707e37 --- /dev/null +++ b/train/tests/test_ui_backend.py @@ -0,0 +1,83 @@ +"""Tests for train UI backend (#101).""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from enpu_train.data.synthetic import make_synthetic_layout_sample +from enpu_train.ui_backend import ( + TrainJobSpec, + inspect_layout_sample, + list_samples_info, + poll_train_job, + start_train_job, + write_train_config, +) + + +def test_inspect_synthetic(tmp_path: Path) -> None: + d = tmp_path / "S001" + make_synthetic_layout_sample(d, sample_id="S001", seed=1) + info = inspect_layout_sample(d) + assert info.ok + assert info.n_systems >= 1 + assert info.n_splits >= 0 + + +def test_list_samples_finds_repo_layout() -> None: + repo_layout = ROOT.parent / "samples" / "layout" + if not repo_layout.is_dir(): + pytest.skip("no samples/layout") + infos = list_samples_info([repo_layout]) + assert any(i.ok for i in infos) or len(infos) >= 0 + + +def test_write_config_and_train_job(tmp_path: Path, monkeypatch) -> None: + for i in range(2): + make_synthetic_layout_sample( + tmp_path / f"S{i}", sample_id=f"S{i}", seed=i, width=320, height=400 + ) + # run under tmp runs + from enpu_train import ui_backend as ub + + monkeypatch.setattr(ub, "DEFAULT_RUNS", tmp_path / "runs") + monkeypatch.setattr(ub, "JOBS_DIR", tmp_path / "runs" / "ui_jobs") + monkeypatch.setattr(ub, "DEFAULT_LAYOUT_ROOT", tmp_path) + + spec = TrainJobSpec( + run_name="test_ui_run", + tasks=["l2", "l3"], + epochs=1, + batch_size=1, + data_roots=[str(tmp_path)], + synth_count=0, + skip_export=True, + device="cpu", + ) + job = start_train_job(spec) + assert job.get("pid") + run_dir = Path(job["run_dir"]) + assert (run_dir / "job.json").is_file() + assert (run_dir / "config.yaml").is_file() + + # wait for process finish (toy train usually < 10s) + import time + + st = {"status": "starting"} + for _ in range(120): + st = poll_train_job(run_dir) + if st["status"] in ("succeeded", "failed"): + break + time.sleep(0.25) + if st["status"] == "failed": + print(st.get("log_tail")) + assert st["status"] == "succeeded", st.get("log_tail") + assert (run_dir / "last.pt").is_file() or (run_dir / "best.pt").is_file() diff --git a/train/ui/app.py b/train/ui/app.py new file mode 100644 index 0000000..08701a1 --- /dev/null +++ b/train/ui/app.py @@ -0,0 +1,419 @@ +"""EnPu Train UI (#101) — Streamlit app. + +Run from train/ directory:: + + streamlit run ui/app.py + +Flow: import .enpu.json -> Layout GT list -> train -> eval -> metrics. +""" + +from __future__ import annotations + +import sys +from datetime import datetime +from pathlib import Path + +import streamlit as st + +TRAIN_ROOT = Path(__file__).resolve().parents[1] +if str(TRAIN_ROOT) not in sys.path: + sys.path.insert(0, str(TRAIN_ROOT)) + +from enpu_train.ui_backend import ( # noqa: E402 + DEFAULT_LAYOUT_ROOT, + DEFAULT_RUNS, + REPO_ROOT, + TrainJobSpec, + cancel_train_job, + cuda_available, + default_ckpt_for_run, + import_enpu_project, + inspect_layout_sample, + list_runs, + list_samples_info, + load_history, + poll_train_job, + run_eval, + start_train_job, +) + +st.set_page_config( + page_title="EnPu Train", + page_icon="🎼", + layout="wide", + initial_sidebar_state="expanded", +) + + +def _init_state() -> None: + ss = st.session_state + ss.setdefault("active_run_dir", None) + ss.setdefault("last_eval", None) + ss.setdefault("import_messages", []) + ss.setdefault("layout_roots", [str(DEFAULT_LAYOUT_ROOT)]) + + +def main() -> None: + _init_state() + st.title("EnPu Train — L1–L3 布局训练") + st.caption( + "父任务 #92 · UI #101 · 数据规范 data-spec · Framework #95。\n" + "本界面**不是**恩谱桌面产品；改框/拖线请在桌面完成后再导入工程。" + ) + + with st.sidebar: + st.header("路径") + st.text(f"REPO: {REPO_ROOT}") + st.text(f"layout: {DEFAULT_LAYOUT_ROOT}") + st.text(f"runs: {DEFAULT_RUNS}") + if cuda_available(): + st.success("CUDA available") + default_device = "cuda" + else: + st.info("CUDA not available — using CPU") + default_device = "cpu" + st.markdown( + """ +**等价 CLI** +```text +python scripts/export_from_enpu_project.py -p song.enpu.json -o ../samples/layout/L00x +python scripts/train.py --config configs/mvp_l2_l3.yaml +python scripts/eval.py --ckpt runs/.../best.pt --data ../samples/layout +``` +""" + ) + + tab_data, tab_train, tab_test, tab_hist = st.tabs( + ["1. 数据集 / 导入", "2. 训练", "3. 测试", "4. 历史实验"] + ) + + # ---------- Dataset ---------- + with tab_data: + st.subheader("从恩谱工程导入 Layout GT") + c1, c2 = st.columns([3, 1]) + with c1: + project_path = st.text_input( + "工程文件路径 (.enpu.json)", + placeholder=r"C:\Users\...\song.enpu.json", + key="project_path", + ) + with c2: + sample_id = st.text_input("样本 ID（可选）", value="", key="sample_id") + + out_name = st.text_input( + "输出目录名（在 samples/layout 下）", + value="", + placeholder="L003_my_song", + key="out_name", + ) + + if st.button("导入 Layout", type="primary", key="btn_import"): + if not project_path.strip(): + st.error("请填写工程路径") + else: + out_dir = None + if out_name.strip(): + out_dir = DEFAULT_LAYOUT_ROOT / out_name.strip() + try: + with st.spinner("导出并校验…"): + result = import_enpu_project( + project_path.strip(), + out_dir=out_dir, + sample_id=sample_id.strip() or None, + ) + st.session_state.import_messages.append(result) + if result["ok"]: + st.success(f"导入成功: {result['out_dir']}") + else: + st.warning(f"已写出但校验未通过: {result['out_dir']}") + st.json(result.get("info")) + except Exception as e: + st.error(f"导入失败: {e}") + + st.divider() + st.subheader("样本列表") + if st.button("刷新列表", key="btn_refresh_samples"): + st.rerun() + + samples = list_samples_info(st.session_state.layout_roots) + if not samples: + st.info("暂无 layout 样本。请导入工程，或确认 samples/layout 存在。") + else: + rows = [] + for s in samples: + rows.append( + { + "id": s.sample_id, + "ok": "✅" if s.ok else "❌", + "systems": s.n_systems, + "splits": s.n_splits, + "measures": s.n_measures, + "size": f"{s.width}x{s.height}" if s.width else "", + "path": s.path, + "errors": "; ".join(s.errors[:2]) if s.errors else "", + } + ) + st.dataframe(rows, use_container_width=True, hide_index=True) + + # preview + ok_samples = [s for s in samples if s.image_path] + if ok_samples: + pick = st.selectbox( + "预览样本", + options=ok_samples, + format_func=lambda s: s.sample_id, + key="preview_sample", + ) + if pick and pick.image_path: + cols = st.columns([1, 1]) + with cols[0]: + st.image(pick.image_path, caption=pick.sample_id, use_container_width=True) + with cols[1]: + st.json(pick.to_dict()) + + # train/val multi-select + st.markdown("**训练用样本根目录**（默认整个 `samples/layout`）") + st.caption("当前 Framework 按目录加载；导入后的样本已在 samples/layout 下即可参与训练。") + use_synth = st.checkbox("训练时附加合成样本", value=True, key="use_synth") + + # ---------- Train ---------- + with tab_train: + st.subheader("一键训练") + tc1, tc2, tc3, tc4 = st.columns(4) + with tc1: + tasks = st.multiselect( + "任务", + options=["l2", "l3"], + default=["l2", "l3"], + key="tasks", + ) + with tc2: + epochs = st.number_input("epochs", min_value=1, max_value=200, value=2, key="epochs") + with tc3: + batch_size = st.number_input("batch size", min_value=1, max_value=32, value=2, key="bs") + with tc4: + device = st.selectbox( + "device", + options=["cpu", "cuda"] if cuda_available() else ["cpu"], + index=0 if default_device == "cpu" else 0, + key="device", + ) + + with st.expander("高级"): + lr = st.number_input("learning rate", value=1e-3, format="%.5f", key="lr") + val_ratio = st.slider("val_ratio", 0.0, 0.5, 0.25, 0.05, key="val_ratio") + synth_count = st.number_input( + "synth_count", + min_value=0, + max_value=64, + value=4 if st.session_state.get("use_synth", True) else 0, + key="synth_count", + ) + run_name = st.text_input( + "实验名 / 输出目录名", + value=f"ui_{datetime.now().strftime('%Y%m%d_%H%M%S')}", + key="run_name", + ) + skip_export = st.checkbox("跳过权重导出（更快）", value=False, key="skip_export") + + b1, b2, b3 = st.columns(3) + with b1: + start = st.button("开始训练", type="primary", key="btn_train") + with b2: + refresh = st.button("刷新进度", key="btn_poll") + with b3: + cancel = st.button("取消训练", key="btn_cancel") + + if start: + if not tasks: + st.error("请至少选择一个任务 (l2/l3)") + else: + spec = TrainJobSpec( + run_name=run_name.strip() or f"ui_{int(datetime.now().timestamp())}", + tasks=list(tasks), + epochs=int(epochs), + batch_size=int(batch_size), + lr=float(lr), + device=str(device), + data_roots=[str(DEFAULT_LAYOUT_ROOT)], + synth_count=int(synth_count), + val_ratio=float(val_ratio), + skip_export=bool(skip_export), + ) + try: + job = start_train_job(spec) + st.session_state.active_run_dir = job["run_dir"] + st.success(f"已启动 PID={job['pid']} → {job['run_dir']}") + except Exception as e: + st.error(f"启动失败: {e}") + + active = st.session_state.active_run_dir + if active is None: + # pick latest run with job + runs = list_runs() + for r in runs: + if (Path(r["path"]) / "job.json").is_file(): + active = r["path"] + break + + if cancel and active: + res = cancel_train_job(active) + if res.get("ok"): + st.warning(f"已请求取消 PID={res.get('pid')}") + else: + st.error(res.get("error") or "取消失败") + + if active and (refresh or start or True): + st.markdown(f"**当前 run:** `{active}`") + status = poll_train_job(active) + st.session_state.active_run_dir = active + col_a, col_b = st.columns(2) + with col_a: + st.metric("状态", status.get("status", "?")) + st.write(f"PID: {status.get('pid')}") + with col_b: + hist = status.get("history") or [] + if hist: + last = hist[-1] + st.metric("epoch", last.get("epoch")) + st.metric("train_loss", f"{last.get('train_loss', float('nan')):.4f}") + val = last.get("val") or {} + if val: + st.write( + f"val L2 IoU={val.get('l2_mean_iou')} · " + f"L3 x_err={val.get('l3_mean_abs_x_error')} · " + f"count_mae={val.get('l3_split_count_mae')}" + ) + + if hist: + chart = { + "train_loss": { + str(h.get("epoch")): h.get("train_loss") + for h in hist + if h.get("train_loss") is not None + } + } + val_loss = { + str(h.get("epoch")): (h.get("val") or {}).get("loss") + for h in hist + if (h.get("val") or {}).get("loss") is not None + } + if val_loss: + chart["val_loss"] = val_loss + try: + st.line_chart(chart) + except Exception: + st.json(hist) + + with st.expander("训练日志 (tail)", expanded=status.get("status") == "running"): + st.code(status.get("log_tail") or "(empty)", language="text") + + if status.get("status") == "running": + st.info("训练进行中 — 点击「刷新进度」更新。") + elif status.get("status") == "failed": + st.error("训练失败，请查看日志。") + elif status.get("status") == "succeeded": + st.success("训练完成。可到「测试」页评估。") + ckpt = default_ckpt_for_run(active) + if ckpt: + st.write(f"ckpt: `{ckpt}`") + + # ---------- Test ---------- + with tab_test: + st.subheader("一键测试 / 评估") + runs = list_runs() + run_options = {r["name"]: r["path"] for r in runs if r.get("has_best") or r.get("has_last")} + if not run_options: + st.warning("还没有可用的 ckpt。请先训练。") + else: + run_pick = st.selectbox("实验", options=list(run_options.keys()), key="eval_run") + run_dir = Path(run_options[run_pick]) + ckpt = default_ckpt_for_run(run_dir) + st.write(f"ckpt: `{ckpt}`") + data_root = st.text_input( + "测试数据目录", + value=str(DEFAULT_LAYOUT_ROOT), + key="eval_data", + ) + eval_device = st.selectbox( + "eval device", + options=["cpu", "cuda"] if cuda_available() else ["cpu"], + key="eval_device", + ) + if st.button("开始测试", type="primary", key="btn_eval"): + if ckpt is None: + st.error("找不到 best.pt / last.pt") + else: + try: + with st.spinner("eval 运行中…"): + res = run_eval( + ckpt, + data_root, + out_json=run_dir / "eval_ui.json", + device=eval_device, + ) + st.session_state.last_eval = res + if res.get("ok"): + st.success("评估完成") + else: + st.error(res.get("error") or "评估失败") + st.code(res.get("stdout") or "") + except Exception as e: + st.error(f"评估异常: {e}") + + last = st.session_state.last_eval + if last and last.get("metrics"): + m = last["metrics"] + st.subheader("指标") + mc1, mc2, mc3, mc4 = st.columns(4) + mc1.metric("L2 mean IoU", _fmt(m.get("l2_mean_iou"))) + mc2.metric("L3 mean |Δx|", _fmt(m.get("l3_mean_abs_x_error"))) + mc3.metric("L3 count MAE", _fmt(m.get("l3_split_count_mae"))) + mc4.metric("L3 count exact", _fmt(m.get("l3_split_count_exact"))) + st.json(m) + elif (run_dir / "eval_ui.json").is_file(): + import json + + m = json.loads((run_dir / "eval_ui.json").read_text(encoding="utf-8")) + st.subheader("上次评估 (eval_ui.json)") + st.json(m) + + # ---------- History ---------- + with tab_hist: + st.subheader("历史实验") + runs = list_runs() + if not runs: + st.info("runs/ 下暂无实验") + else: + for r in runs: + with st.expander( + f"{r['name']} · {r.get('mtime', '')[:19]} · " + f"{'best' if r.get('has_best') else '—'} · " + f"loss={r.get('last_train_loss')}", + expanded=False, + ): + st.write(r["path"]) + if r.get("last_val"): + st.json(r["last_val"]) + hist = load_history(r["path"]) + if hist: + st.write(f"epochs recorded: {len(hist)}") + if st.button("设为当前 run", key=f"set_{r['name']}"): + st.session_state.active_run_dir = r["path"] + st.success(f"active → {r['path']}") + + +def _fmt(v) -> str: + try: + if v is None: + return "—" + x = float(v) + if x != x: + return "nan" + return f"{x:.4f}" + except Exception: + return str(v) + + +if __name__ == "__main__": + main()