From 0737ca3c0db50aad5baf0904d39c9c9b277e34d7 Mon Sep 17 00:00:00 2001 From: loootte <46289941+loootte@users.noreply.github.com> Date: Sun, 2 Aug 2026 18:36:22 +0800 Subject: [PATCH] feat(core): load L1-L3 train weights into structure pipeline (#104) Add structure/learned loader+infer, ENPU_STRUCTURE_L1L3_ENGINE switch with rule fallback, eval_l1l3_engines comparison script, and core-inference docs. Closes #104 --- core/README.md | 7 +- core/app/config.py | 8 + .../pipeline/structure/learned/__init__.py | 12 + .../app/pipeline/structure/learned/adapter.py | 96 ++++++++ .../pipeline/structure/learned/infer_l1l3.py | 175 ++++++++++++++ core/app/pipeline/structure/learned/loader.py | 120 ++++++++++ core/app/pipeline/structure/learned/model.py | 117 ++++++++++ .../pipeline/structure/learned/postprocess.py | 129 +++++++++++ core/app/pipeline/structure/pipeline.py | 73 +++++- core/requirements.txt | 3 + core/tests/test_learned_l1l3.py | 166 ++++++++++++++ docs/train/core-inference.md | 92 ++++++++ scripts/eval_l1l3_engines.py | 217 ++++++++++++++++++ 13 files changed, 1205 insertions(+), 10 deletions(-) create mode 100644 core/app/pipeline/structure/learned/__init__.py create mode 100644 core/app/pipeline/structure/learned/adapter.py create mode 100644 core/app/pipeline/structure/learned/infer_l1l3.py create mode 100644 core/app/pipeline/structure/learned/loader.py create mode 100644 core/app/pipeline/structure/learned/model.py create mode 100644 core/app/pipeline/structure/learned/postprocess.py create mode 100644 core/tests/test_learned_l1l3.py create mode 100644 docs/train/core-inference.md create mode 100644 scripts/eval_l1l3_engines.py diff --git a/core/README.md b/core/README.md index 253c3e6..1d5236d 100644 --- a/core/README.md +++ b/core/README.md @@ -96,6 +96,10 @@ ENPU_RECOGNIZE_ENGINE=mock pytest -q | `ENPU_MAX_UPLOAD_BYTES` | `20971520` | 最大上传 | | `ENPU_CORS_ORIGINS` | `*` | CORS | | `ENPU_PIPELINE_MODE` | `legacy` | `legacy`：整页 OCR→parse；`structure`：L1–L5 结构优先（#58） | +| `ENPU_STRUCTURE_L1L3_ENGINE` | `rule` | `rule` \| `learned`（#104；learned 需 torch + 权重） | +| `ENPU_L1L3_WEIGHTS` | | `best.pt` / `layout_net.pt` 路径 | +| `ENPU_L1L3_DEVICE` | `cpu` | `cpu` / `cuda` | +| `ENPU_L1L3_FALLBACK` | `rule` | 失败时回退 rule 或 `none` | ## 流水线说明 @@ -135,7 +139,8 @@ python scripts/export_layout_gt.py --project path\to\song.enpu.json --out sample ``` 模块：`app/layout_gt/`（export + validate）。 -规范：[l1-l3-data-spec.md](../docs/train/l1-l3-data-spec.md) · 模型方案：[l1-l3-model-design.md](../docs/train/l1-l3-model-design.md)（#94）。 +规范：[l1-l3-data-spec.md](../docs/train/l1-l3-data-spec.md) · 模型方案：[l1-l3-model-design.md](../docs/train/l1-l3-model-design.md)（#94）。 +推理接入：[core-inference.md](../docs/train/core-inference.md)（#104：`structure/learned/`）。 ## Sidecar 打包（可选，Issue #8） diff --git a/core/app/config.py b/core/app/config.py index aa51388..5d448c4 100644 --- a/core/app/config.py +++ b/core/app/config.py @@ -27,6 +27,14 @@ class Settings(BaseSettings): # legacy = OCR-first (current default); structure = L1–L5 geometry-first (#58) pipeline_mode: str = "legacy" + # structure L1–L3 engine (#104): rule (OpenCV) | learned (train weights) + structure_l1l3_engine: str = "rule" + # path to layout_net.pt / best.pt (required when engine=learned) + l1l3_weights: str = "" + l1l3_device: str = "cpu" + # on load/infer failure: rule | none (raise) + l1l3_fallback: str = "rule" + # OCR / preprocess ocr_lang: str = "ch" ocr_use_angle_cls: bool = True diff --git a/core/app/pipeline/structure/learned/__init__.py b/core/app/pipeline/structure/learned/__init__.py new file mode 100644 index 0000000..bbd0bd7 --- /dev/null +++ b/core/app/pipeline/structure/learned/__init__.py @@ -0,0 +1,12 @@ +"""Learned L1–L3 layout inference for structure pipeline (#104). + +Does **not** import the train/ app. Weights format: ``enpu_layout_net_v0`` +or train ``best.pt`` / ``last.pt`` (state_dict + cfg). +""" + +from app.pipeline.structure.learned.infer_l1l3 import ( + LearnedL1L3Error, + run_learned_l2_l3, +) + +__all__ = ["LearnedL1L3Error", "run_learned_l2_l3"] diff --git a/core/app/pipeline/structure/learned/adapter.py b/core/app/pipeline/structure/learned/adapter.py new file mode 100644 index 0000000..0903911 --- /dev/null +++ b/core/app/pipeline/structure/learned/adapter.py @@ -0,0 +1,96 @@ +"""Convert learned systems + splits → PageLayout skeleton (#104).""" + +from __future__ import annotations + +from typing import Any + +from app.pipeline.structure.ir import ( + PageLayout, + PageRegion, + Rect, + RegionRole, + SplitLine, + StaffSystem, +) +from app.pipeline.structure.splits import normalize_splits, splits_to_measures + + +def systems_splits_to_page_layout( + *, + width: int, + height: int, + system_boxes: list[dict[str, float]], + splits_per_system: list[list[float]], + score_region: Rect | None = None, + title_box: Rect | None = None, + key_time_box: Rect | None = None, + warnings: list[str] | None = None, + engine_meta: dict[str, Any] | None = None, + min_gap: float = 6.0, + min_measure_width: float = 4.0, +) -> PageLayout: + """Build PageLayout with L1 regions + L2 systems + L3 splits/measures.""" + regions: list[PageRegion] = [] + if title_box is not None: + regions.append( + PageRegion(role=RegionRole.title, rect=title_box, confidence=0.6) + ) + if key_time_box is not None: + regions.append( + PageRegion(role=RegionRole.key_time, rect=key_time_box, confidence=0.55) + ) + if score_region is None: + score_region = Rect(0, 0, float(width), float(height)) + regions.append( + PageRegion(role=RegionRole.score, rect=score_region, confidence=0.75) + ) + + systems: list[StaffSystem] = [] + for i, box in enumerate(system_boxes): + rect = Rect( + float(box["x1"]), + float(box["y1"]), + float(box["x2"]), + float(box["y2"]), + ) + raw_xs = splits_per_system[i] if i < len(splits_per_system) else [] + raw_splits = [ + SplitLine(x=float(x), split_id=f"s{i}-{j}", source="detect") + for j, x in enumerate(raw_xs) + ] + splits = normalize_splits( + raw_splits, + x_left=rect.x1, + x_right=rect.x2, + min_gap=min_gap, + default_source="detect", + ) + measures = splits_to_measures( + x_left=rect.x1, + x_right=rect.x2, + y_top=rect.y1, + y_bot=rect.y2, + splits=splits, + min_measure_width=min_measure_width, + measure_source="l3_split", + ) + systems.append( + StaffSystem( + index=i, + rect=rect, + measures=measures, + barline_xs=[s.x for s in splits], + splits=splits, + confidence=0.7, + extra={"source": "learned_l1l3", "engine": "learned"}, + ) + ) + + return PageLayout( + width=width, + height=height, + regions=regions, + systems=systems, + warnings=list(warnings or []), + debug={"l1l3_engine": "learned", **(engine_meta or {})}, + ) diff --git a/core/app/pipeline/structure/learned/infer_l1l3.py b/core/app/pipeline/structure/learned/infer_l1l3.py new file mode 100644 index 0000000..a90f742 --- /dev/null +++ b/core/app/pipeline/structure/learned/infer_l1l3.py @@ -0,0 +1,175 @@ +"""Run learned L2+L3 on a BGR image → PageLayout skeleton (#104). + +L1 uses optional rule regions or full-page score ROI (MVP: hybrid). +""" + +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np + +from app.config import Settings +from app.pipeline.structure.ir import PageLayout, PageRegion, Rect, RegionRole +from app.pipeline.structure.learned.adapter import systems_splits_to_page_layout +from app.pipeline.structure.learned.loader import ( + WeightsLoadError, + get_cached_layout_model, +) +from app.pipeline.structure.learned.postprocess import ( + bgr_to_model_tensor, + l2_heat_to_system_boxes, + l3_heat_to_split_xs, +) + +logger = logging.getLogger(__name__) + + +class LearnedL1L3Error(Exception): + """Inference failure (caller may fallback to rule).""" + + +def run_learned_l2_l3( + image_bgr: np.ndarray, + *, + settings: Settings, + l1_regions: list[PageRegion] | None = None, +) -> PageLayout: + """Infer L2 systems + L3 splits; attach L1 regions if provided.""" + if image_bgr is None or image_bgr.size == 0: + raise LearnedL1L3Error("empty image") + + weights = (settings.l1l3_weights or "").strip() + if not weights: + raise LearnedL1L3Error( + "ENPU_L1L3_WEIGHTS is empty; set path to train export " + "(layout_net.pt / best.pt)" + ) + + device = (settings.l1l3_device or "cpu").strip() or "cpu" + try: + model, wmeta = get_cached_layout_model(weights, device) + except WeightsLoadError as e: + raise LearnedL1L3Error(str(e)) from e + + try: + import torch + except ImportError as e: + raise LearnedL1L3Error("torch not installed") from e + + h, w = image_bgr.shape[:2] + cfg = model.cfg + page_h, page_w = cfg.page_h, cfg.page_w + row_h, row_w = cfg.row_h, cfg.row_w + + page_t = bgr_to_model_tensor(image_bgr, out_h=page_h, out_w=page_w).to(device) + warnings: list[str] = [ + "l1l3_engine=learned (#104)", + f"l1l3_weights={wmeta.get('path')}", + f"l1l3_format={wmeta.get('format')}", + ] + + with torch.no_grad(): + out = model(page=page_t) + if "l2_logits" not in out: + raise LearnedL1L3Error("model has no L2 head / tasks") + l2_prob = torch.sigmoid(out["l2_logits"])[0].detach().cpu().numpy() + + system_boxes = l2_heat_to_system_boxes( + l2_prob, + orig_h=h, + orig_w=w, + page_h=page_h, + page_w=page_w, + ) + if not system_boxes: + raise LearnedL1L3Error("L2 produced no systems") + + # L3 per system crop + splits_per: list[list[float]] = [[] for _ in system_boxes] + row_tensors: list = [] + row_indices: list[int] = [] + crop_meta: list[dict[str, float]] = [] + + for i, box in enumerate(system_boxes): + x1 = int(max(0, box["x1"])) + y1 = int(max(0, box["y1"])) + x2 = int(min(w, box["x2"])) + y2 = int(min(h, box["y2"])) + if x2 <= x1 + 2 or y2 <= y1 + 2: + continue + pad_x = int(0.02 * (x2 - x1)) + pad_y = int(0.05 * (y2 - y1)) + cx1 = max(0, x1 - pad_x) + cy1 = max(0, y1 - pad_y) + cx2 = min(w, x2 + pad_x) + cy2 = min(h, y2 + pad_y) + crop = image_bgr[cy1:cy2, cx1:cx2] + row_tensors.append(bgr_to_model_tensor(crop, out_h=row_h, out_w=row_w)[0]) + row_indices.append(i) + crop_meta.append( + { + "x_left": float(cx1), + "x_right": float(cx2), + "y1": float(cy1), + "y2": float(cy2), + } + ) + + if row_tensors and "l3" in cfg.tasks: + rows_batch = torch.stack(row_tensors, dim=0).to(device) + with torch.no_grad(): + l3_out = model(rows=rows_batch) + l3_logits = l3_out.get("l3_logits") + if l3_logits is not None: + probs = torch.sigmoid(l3_logits).detach().cpu().numpy() + for j, meta in enumerate(crop_meta): + xs = l3_heat_to_split_xs( + probs[j], + x_left=meta["x_left"], + x_right=meta["x_right"], + ) + splits_per[row_indices[j]] = xs + else: + warnings.append("L3 skipped (no crops or task disabled)") + + # L1 regions + score_rect = None + title_box = None + key_time_box = None + if l1_regions: + for r in l1_regions: + if r.role == RegionRole.score: + score_rect = r.rect + elif r.role == RegionRole.title: + title_box = r.rect + elif r.role == RegionRole.key_time: + key_time_box = r.rect + if score_rect is None: + # score band covering all systems + y1 = min(b["y1"] for b in system_boxes) + y2 = max(b["y2"] for b in system_boxes) + score_rect = Rect(0, max(0.0, y1 - 0.02 * h), float(w), min(float(h), y2 + 0.02 * h)) + + warnings.append( + f"L2: learned {len(system_boxes)} system(s); " + f"L3 splits total={sum(len(s) for s in splits_per)}" + ) + + layout = systems_splits_to_page_layout( + width=w, + height=h, + system_boxes=system_boxes, + splits_per_system=splits_per, + score_region=score_rect, + title_box=title_box, + key_time_box=key_time_box, + warnings=warnings, + engine_meta={ + "weights": wmeta, + "n_systems": len(system_boxes), + "n_splits": sum(len(s) for s in splits_per), + }, + ) + return layout diff --git a/core/app/pipeline/structure/learned/loader.py b/core/app/pipeline/structure/learned/loader.py new file mode 100644 index 0000000..a25bebe --- /dev/null +++ b/core/app/pipeline/structure/learned/loader.py @@ -0,0 +1,120 @@ +"""Load EnPu layout weights (#104).""" + +from __future__ import annotations + +import logging +from functools import lru_cache +from pathlib import Path +from typing import Any + +from app.pipeline.structure.learned.model import LayoutNet, LayoutNetConfig + +logger = logging.getLogger(__name__) + +WEIGHTS_FORMAT = "enpu_layout_net_v0" + + +class WeightsLoadError(Exception): + """Missing / corrupt / incompatible weights.""" + + +def _require_torch(): + try: + import torch + except ImportError as e: + raise WeightsLoadError( + "torch is required for ENPU_STRUCTURE_L1L3_ENGINE=learned. " + "Install with: pip install torch (optional; default engine remains rule)" + ) from e + return torch + + +def _cfg_from_payload(payload: dict[str, Any]) -> LayoutNetConfig: + # Prefer export layout_net_v0 fields; fall back to train TrainConfig blob + cfg_blob = payload.get("cfg") or {} + l2_heat = int( + payload.get("l2_heat_len") + or cfg_blob.get("l2_heat_len") + or 128 + ) + l3_heat = int( + payload.get("l3_heat_len") + or cfg_blob.get("l3_heat_len") + or 128 + ) + tasks_raw = payload.get("tasks") or cfg_blob.get("tasks") or ("l2", "l3") + tasks = tuple(str(t) for t in tasks_raw) + return LayoutNetConfig( + l2_heat_len=l2_heat, + l3_heat_len=l3_heat, + tasks=tasks if tasks else ("l2", "l3"), + page_h=int(cfg_blob.get("page_h") or 384), + page_w=int(cfg_blob.get("page_w") or 512), + row_h=int(cfg_blob.get("row_h") or 64), + row_w=int(cfg_blob.get("row_w") or 256), + base_channels=16, + ) + + +def load_layout_weights( + path: str | Path, + *, + device: str = "cpu", +) -> tuple[LayoutNet, dict[str, Any]]: + """Load ``enpu_layout_net_v0`` or train ``best.pt`` / ``last.pt``. + + Returns (model.eval(), meta dict with format/path/cfg fields). + """ + torch = _require_torch() + path = Path(path) + if not path.is_file(): + raise WeightsLoadError(f"weights not found: {path}") + + try: + payload = torch.load(path, map_location=device, weights_only=False) + except Exception as e: + raise WeightsLoadError(f"failed to load weights {path}: {e}") from e + + if not isinstance(payload, dict) or "model" not in payload: + raise WeightsLoadError( + f"unsupported weights file (need dict with 'model' state_dict): {path}" + ) + + fmt = str(payload.get("format") or "train_ckpt") + cfg = _cfg_from_payload(payload) + model = LayoutNet(cfg) + try: + model.load_state_dict(payload["model"], strict=True) + except Exception as e: + raise WeightsLoadError(f"state_dict mismatch for {path}: {e}") from e + + model.to(device) + model.eval() + meta = { + "format": fmt, + "path": str(path.resolve()), + "tasks": list(cfg.tasks), + "l2_heat_len": cfg.l2_heat_len, + "l3_heat_len": cfg.l3_heat_len, + "page_h": cfg.page_h, + "page_w": cfg.page_w, + "row_h": cfg.row_h, + "row_w": cfg.row_w, + "device": device, + "weights_format_expected": WEIGHTS_FORMAT, + } + logger.info("loaded layout weights %s format=%s tasks=%s", path, fmt, cfg.tasks) + return model, meta + + +@lru_cache(maxsize=4) +def get_cached_layout_model( + path: str, + device: str = "cpu", +) -> tuple[LayoutNet, dict[str, Any]]: + """Process-level cache keyed by path+device.""" + return load_layout_weights(path, device=device) + + +def clear_layout_model_cache() -> None: + get_cached_layout_model.cache_clear() diff --git a/core/app/pipeline/structure/learned/model.py b/core/app/pipeline/structure/learned/model.py new file mode 100644 index 0000000..92a61e5 --- /dev/null +++ b/core/app/pipeline/structure/learned/model.py @@ -0,0 +1,117 @@ +"""Minimal LayoutNet (L2 page y-heat + L3 row x-heat) for core inference (#104). + +Architecture mirrors ``train/enpu_train/models/layout_net.py`` so exported +weights load without importing the train package. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class ConvBNReLU(nn.Module): + def __init__(self, c_in: int, c_out: int, k: int = 3, s: int = 1) -> None: + super().__init__() + p = k // 2 + self.net = nn.Sequential( + nn.Conv2d(c_in, c_out, k, stride=s, padding=p, bias=False), + nn.BatchNorm2d(c_out), + nn.ReLU(inplace=True), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class PageL2Head(nn.Module): + def __init__(self, heat_len: int = 128, base: int = 16) -> None: + super().__init__() + self.heat_len = heat_len + self.enc = nn.Sequential( + ConvBNReLU(3, base, 3, 2), + ConvBNReLU(base, base * 2, 3, 2), + ConvBNReLU(base * 2, base * 4, 3, 2), + ConvBNReLU(base * 4, base * 4, 3, 2), + ) + self.head = nn.Sequential( + nn.Conv1d(base * 4, base * 2, 3, padding=1), + nn.ReLU(inplace=True), + nn.Conv1d(base * 2, 1, 1), + ) + + def forward(self, page: torch.Tensor) -> torch.Tensor: + f = self.enc(page) + f = f.mean(dim=3) + f = F.interpolate( + f.unsqueeze(-1), + size=(self.heat_len, 1), + mode="bilinear", + align_corners=False, + ).squeeze(-1) + return self.head(f).squeeze(1) + + +class RowL3Head(nn.Module): + def __init__(self, heat_len: int = 128, base: int = 16) -> None: + super().__init__() + self.heat_len = heat_len + self.enc = nn.Sequential( + ConvBNReLU(3, base, 3, 2), + ConvBNReLU(base, base * 2, 3, 2), + ConvBNReLU(base * 2, base * 4, 3, 2), + ConvBNReLU(base * 4, base * 4, 3, 2), + ) + self.head = nn.Sequential( + nn.Conv1d(base * 4, base * 2, 3, padding=1), + nn.ReLU(inplace=True), + nn.Conv1d(base * 2, 1, 1), + ) + + def forward(self, rows: torch.Tensor) -> torch.Tensor: + if rows.numel() == 0: + return rows.new_zeros((0, self.heat_len)) + f = self.enc(rows) + f = f.mean(dim=2) + f = F.interpolate( + f.unsqueeze(2), + size=(1, self.heat_len), + mode="bilinear", + align_corners=False, + ).squeeze(2) + return self.head(f).squeeze(1) + + +@dataclass +class LayoutNetConfig: + l2_heat_len: int = 128 + l3_heat_len: int = 128 + base_channels: int = 16 + tasks: tuple[str, ...] = ("l2", "l3") + page_h: int = 384 + page_w: int = 512 + row_h: int = 64 + row_w: int = 256 + + +class LayoutNet(nn.Module): + def __init__(self, cfg: LayoutNetConfig | None = None) -> None: + super().__init__() + self.cfg = cfg or LayoutNetConfig() + self.l2 = PageL2Head(self.cfg.l2_heat_len, self.cfg.base_channels) + self.l3 = RowL3Head(self.cfg.l3_heat_len, self.cfg.base_channels) + + def forward( + self, + page: torch.Tensor | None = None, + rows: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + out: dict[str, torch.Tensor] = {} + if page is not None and "l2" in self.cfg.tasks: + out["l2_logits"] = self.l2(page) + if rows is not None and "l3" in self.cfg.tasks: + out["l3_logits"] = self.l3(rows) + return out diff --git a/core/app/pipeline/structure/learned/postprocess.py b/core/app/pipeline/structure/learned/postprocess.py new file mode 100644 index 0000000..497d2a1 --- /dev/null +++ b/core/app/pipeline/structure/learned/postprocess.py @@ -0,0 +1,129 @@ +"""Decode heatmaps → system bands + interior split xs (#104).""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + + +def decode_peaks( + heat: np.ndarray, + *, + min_prominence: float = 0.25, + min_gap: int = 4, +) -> list[int]: + h = np.asarray(heat, dtype=np.float32).reshape(-1) + peaks: list[tuple[float, int]] = [] + for i in range(1, len(h) - 1): + if h[i] >= min_prominence and h[i] >= h[i - 1] and h[i] >= h[i + 1]: + peaks.append((float(h[i]), i)) + peaks.sort(reverse=True) + chosen: list[int] = [] + for _, i in peaks: + if all(abs(i - j) >= min_gap for j in chosen): + chosen.append(i) + return sorted(chosen) + + +def l2_heat_to_system_boxes( + heat: np.ndarray, + *, + orig_h: int, + orig_w: int, + page_h: int, + page_w: int, + band_frac: float = 0.09, + min_prominence: float = 0.22, + x_margin_frac: float = 0.02, +) -> list[dict[str, float]]: + """Map L2 y-heat peaks to full-width-ish horizontal system bboxes (orig pixels).""" + heat = np.asarray(heat, dtype=np.float32).reshape(-1) + gap = max(3, len(heat) // 40) + peaks = decode_peaks(heat, min_prominence=min_prominence, min_gap=gap) + if not peaks: + # fallback: single mid band + cy = orig_h * 0.5 + half = orig_h * band_frac * 0.5 + xm = orig_w * x_margin_frac + return [ + { + "x1": xm, + "y1": max(0.0, cy - half), + "x2": float(orig_w) - xm, + "y2": min(float(orig_h), cy + half), + } + ] + + # estimate band height from peak spacing + if len(peaks) >= 2: + spacings = [ + (peaks[i + 1] - peaks[i]) / max(1, len(heat) - 1) * orig_h + for i in range(len(peaks) - 1) + ] + half = 0.35 * float(np.median(spacings)) + else: + half = orig_h * band_frac * 0.5 + half = max(half, orig_h * 0.03) + half = min(half, orig_h * 0.12) + + xm = orig_w * x_margin_frac + boxes: list[dict[str, float]] = [] + for p in peaks: + cy = p / max(1, len(heat) - 1) * orig_h + boxes.append( + { + "x1": float(xm), + "y1": float(max(0.0, cy - half)), + "x2": float(orig_w - xm), + "y2": float(min(float(orig_h), cy + half)), + } + ) + # sort top→bottom, drop heavy overlaps + boxes.sort(key=lambda b: b["y1"]) + cleaned: list[dict[str, float]] = [] + for b in boxes: + if cleaned and b["y1"] < cleaned[-1]["y2"] - 0.3 * (cleaned[-1]["y2"] - cleaned[-1]["y1"]): + # merge into previous by expanding + cleaned[-1]["y2"] = max(cleaned[-1]["y2"], b["y2"]) + cleaned[-1]["y1"] = min(cleaned[-1]["y1"], b["y1"]) + else: + cleaned.append(b) + return cleaned + + +def l3_heat_to_split_xs( + heat: np.ndarray, + *, + x_left: float, + x_right: float, + min_prominence: float = 0.28, +) -> list[float]: + """Map L3 x-heat peaks to full-image interior split x positions.""" + heat = np.asarray(heat, dtype=np.float32).reshape(-1) + gap = max(3, len(heat) // 32) + peaks = decode_peaks(heat, min_prominence=min_prominence, min_gap=gap) + width = max(1e-3, x_right - x_left) + xs: list[float] = [] + for p in peaks: + rel = p / max(1, len(heat) - 1) + x = x_left + rel * width + if x_left + 1.0 < x < x_right - 1.0: + xs.append(float(x)) + return sorted(xs) + + +def bgr_to_model_tensor( + bgr: Any, + *, + out_h: int, + out_w: int, +) -> Any: + """Resize BGR uint8 → float CHW tensor in [0,1] RGB order (torch).""" + import cv2 + import torch + + rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) + resized = cv2.resize(rgb, (out_w, out_h), interpolation=cv2.INTER_LINEAR) + t = torch.from_numpy(np.ascontiguousarray(resized)).float() / 255.0 + return t.permute(2, 0, 1).unsqueeze(0) # 1,3,H,W diff --git a/core/app/pipeline/structure/pipeline.py b/core/app/pipeline/structure/pipeline.py index 19107cd..9f1d6ad 100644 --- a/core/app/pipeline/structure/pipeline.py +++ b/core/app/pipeline/structure/pipeline.py @@ -110,20 +110,75 @@ def run_structure_recognize( if geometric: warnings.append(f"preprocess_toolbox: {' → '.join(pre.steps)}") - # --- L1 + l1l3_engine = (settings.structure_l1l3_engine or "rule").strip().lower() + used_engine = "rule" + layout_debug: dict = {"preprocess_steps": list(pre.steps), "scale": pre.scale} + + # --- L1 (always rule geometry for score/title; learned may refine score band) regions, w1 = detect_page_regions(work) warnings.extend(w1) score_rect = next((r.rect for r in regions if r.role == RegionRole.score), None) if score_rect is None: score_rect = Rect(0, 0, float(w), float(h)) - # --- L2 - systems, w2 = detect_staff_systems(work, score_rect) - warnings.extend(w2) + systems = [] + if l1l3_engine == "learned": + try: + from app.pipeline.structure.learned.infer_l1l3 import ( + LearnedL1L3Error, + run_learned_l2_l3, + ) + + learned_layout = run_learned_l2_l3( + work, settings=settings, l1_regions=regions + ) + # Prefer learned systems + measures/splits; keep rule L1 if richer + systems = learned_layout.systems + if learned_layout.regions: + # merge: keep rule title/key_time, use learned score if present + rule_by_role = {r.role: r for r in regions} + learned_by_role = {r.role: r for r in learned_layout.regions} + merged = [] + for role in (RegionRole.title, RegionRole.key_time, RegionRole.score): + if role in rule_by_role and role != RegionRole.score: + merged.append(rule_by_role[role]) + elif role in learned_by_role: + merged.append(learned_by_role[role]) + elif role in rule_by_role: + merged.append(rule_by_role[role]) + regions = merged or regions + score_rect = next( + (r.rect for r in regions if r.role == RegionRole.score), + score_rect, + ) + warnings.extend(learned_layout.warnings) + layout_debug["l1l3"] = learned_layout.debug + used_engine = "learned" + warnings.append("l1l3_engine=learned (L2+L3 from weights; L1 hybrid)") + except Exception as exc: # noqa: BLE001 — fallback path + fb = (settings.l1l3_fallback or "rule").strip().lower() + msg = f"learned L1–L3 failed: {exc}" + logger.warning(msg) + warnings.append(msg) + if fb == "rule": + warnings.append("l1l3_fallback=rule (#104)") + used_engine = "rule" + systems, w2 = detect_staff_systems(work, score_rect) + warnings.extend(w2) + systems, w3 = segment_measures_on_systems(work, systems) + warnings.extend(w3) + else: + raise StructurePipelineError(msg, status_code=500) from exc + else: + warnings.append("l1l3_engine=rule") + # --- L2 + systems, w2 = detect_staff_systems(work, score_rect) + warnings.extend(w2) + # --- L3 + systems, w3 = segment_measures_on_systems(work, systems) + warnings.extend(w3) - # --- L3 - systems, w3 = segment_measures_on_systems(work, systems) - warnings.extend(w3) + layout_debug["l1l3_engine"] = used_engine # --- L4 systems, w4 = detect_note_candidates(work, systems) @@ -160,7 +215,7 @@ def run_structure_recognize( time_signature=time_sig, title=title, warnings=warnings, - debug={"preprocess_steps": list(pre.steps), "scale": pre.scale}, + debug=layout_debug, ) return _layout_to_response( @@ -169,7 +224,7 @@ def run_structure_recognize( filename=filename, content_type=content_type, started=started, - preprocess_steps=list(pre.steps) + ["structure_l1_l5"], + preprocess_steps=list(pre.steps) + ["structure_l1_l5", f"l1l3={used_engine}"], ) diff --git a/core/requirements.txt b/core/requirements.txt index 8a78571..9746d18 100644 --- a/core/requirements.txt +++ b/core/requirements.txt @@ -25,6 +25,9 @@ music21>=9.1,<10 # Layer param YAML (#89) PyYAML>=6.0,<7 +# Optional: learned L1–L3 structure engine (#104). Not required for rule path / CI / slim installer. +# pip install torch + # tests httpx>=0.27,<1 pytest>=8.0,<9 diff --git a/core/tests/test_learned_l1l3.py b/core/tests/test_learned_l1l3.py new file mode 100644 index 0000000..7fb68d7 --- /dev/null +++ b/core/tests/test_learned_l1l3.py @@ -0,0 +1,166 @@ +"""Tests for learned L1–L3 structure engine (#104).""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +from app.config import Settings, clear_settings_cache +from app.pipeline.structure.learned.adapter import systems_splits_to_page_layout +from app.pipeline.structure.learned.loader import ( + WeightsLoadError, + clear_layout_model_cache, + load_layout_weights, +) +from app.pipeline.structure.learned.postprocess import ( + decode_peaks, + l2_heat_to_system_boxes, + l3_heat_to_split_xs, +) +from app.pipeline.structure.pipeline import run_structure_recognize + + +def test_decode_peaks_and_boxes() -> None: + heat = np.zeros(64, dtype=np.float32) + heat[10] = 0.9 + heat[40] = 0.8 + peaks = decode_peaks(heat, min_prominence=0.3, min_gap=5) + assert 10 in peaks and 40 in peaks + boxes = l2_heat_to_system_boxes( + heat, orig_h=400, orig_w=300, page_h=384, page_w=512 + ) + assert len(boxes) >= 1 + xs = l3_heat_to_split_xs(heat, x_left=10, x_right=290) + assert all(10 < x < 290 for x in xs) + + +def test_adapter_normalize_splits() -> None: + layout = systems_splits_to_page_layout( + width=400, + height=300, + system_boxes=[{"x1": 20, "y1": 80, "x2": 380, "y2": 140}], + splits_per_system=[[100, 200, 300]], + ) + assert len(layout.systems) == 1 + sys = layout.systems[0] + assert len(sys.splits) == 3 + assert len(sys.measures) == 4 + assert sys.measures[0].rect.x1 == pytest.approx(20) + + +def test_load_missing_weights() -> None: + with pytest.raises(WeightsLoadError): + load_layout_weights(Path("does_not_exist.pt")) + + +def test_load_real_weights_if_present() -> None: + root = Path(__file__).resolve().parents[2] + candidates = [ + root / "train" / "runs" / "mvp_l2_l3" / "export" / "layout_net.pt", + root / "train" / "runs" / "mvp_l2_l3" / "best.pt", + ] + path = next((p for p in candidates if p.is_file()), None) + if path is None: + pytest.skip("no train weights in repo") + clear_layout_model_cache() + model, meta = load_layout_weights(path, device="cpu") + assert "l2" in meta["tasks"] + # forward smoke + import torch + + page = torch.rand(1, 3, model.cfg.page_h, model.cfg.page_w) + with torch.no_grad(): + out = model(page=page) + assert out["l2_logits"].shape[-1] == model.cfg.l2_heat_len + + +def test_pipeline_rule_default_unchanged(monkeypatch: pytest.MonkeyPatch) -> None: + clear_settings_cache() + monkeypatch.setenv("ENPU_RECOGNIZE_ENGINE", "mock") + monkeypatch.setenv("ENPU_PIPELINE_MODE", "structure") + monkeypatch.setenv("ENPU_STRUCTURE_L1L3_ENGINE", "rule") + clear_settings_cache() + # synthetic-ish page + img = np.full((200, 300, 3), 255, dtype=np.uint8) + img[60:90, 20:280] = 30 + img[120:150, 20:280] = 30 + for x in (40, 100, 160, 220): + img[55:155, x : x + 2] = 0 + import cv2 + + ok, buf = cv2.imencode(".png", img) + assert ok + settings = Settings() + resp = run_structure_recognize(buf.tobytes(), settings=settings, filename="t.png") + assert resp.ok + assert resp.structure is not None + assert any("l1l3_engine=rule" in w for w in (resp.meta.parse_warnings or [])) + + +def test_pipeline_learned_fallback_when_no_weights( + monkeypatch: pytest.MonkeyPatch, +) -> None: + clear_settings_cache() + clear_layout_model_cache() + monkeypatch.setenv("ENPU_RECOGNIZE_ENGINE", "mock") + monkeypatch.setenv("ENPU_PIPELINE_MODE", "structure") + monkeypatch.setenv("ENPU_STRUCTURE_L1L3_ENGINE", "learned") + monkeypatch.setenv("ENPU_L1L3_WEIGHTS", "") + monkeypatch.setenv("ENPU_L1L3_FALLBACK", "rule") + clear_settings_cache() + img = np.full((200, 300, 3), 255, dtype=np.uint8) + img[60:90, 20:280] = 30 + import cv2 + + ok, buf = cv2.imencode(".png", img) + assert ok + settings = Settings() + resp = run_structure_recognize(buf.tobytes(), settings=settings, filename="t.png") + assert resp.ok + warns = resp.meta.parse_warnings or [] + assert any("fallback" in w.lower() or "failed" in w.lower() for w in warns) + + +def test_pipeline_learned_with_weights_if_present( + monkeypatch: pytest.MonkeyPatch, +) -> None: + root = Path(__file__).resolve().parents[2] + path = root / "train" / "runs" / "mvp_l2_l3" / "best.pt" + if not path.is_file(): + pytest.skip("no weights") + clear_settings_cache() + clear_layout_model_cache() + monkeypatch.setenv("ENPU_RECOGNIZE_ENGINE", "mock") + monkeypatch.setenv("ENPU_PIPELINE_MODE", "structure") + monkeypatch.setenv("ENPU_STRUCTURE_L1L3_ENGINE", "learned") + monkeypatch.setenv("ENPU_L1L3_WEIGHTS", str(path)) + monkeypatch.setenv("ENPU_L1L3_DEVICE", "cpu") + monkeypatch.setenv("ENPU_L1L3_FALLBACK", "rule") + clear_settings_cache() + + # use real layout sample if available + sample = root / "samples" / "layout" / "L001_zuozai_baozuo" / "image.png" + if sample.is_file(): + data = sample.read_bytes() + else: + img = np.full((400, 600, 3), 255, dtype=np.uint8) + for y0 in (80, 160, 240, 320): + img[y0 : y0 + 40, 30:570] = 40 + for x in range(80, 560, 90): + img[y0 : y0 + 40, x : x + 3] = 255 + import cv2 + + ok, buf = cv2.imencode(".png", img) + assert ok + data = buf.tobytes() + + settings = Settings() + resp = run_structure_recognize(data, settings=settings, filename="learned.png") + assert resp.ok + assert resp.structure is not None + assert resp.structure.items + warns = " ".join(resp.meta.parse_warnings or []) + # either learned succeeded or fell back + assert "l1l3" in warns.lower() or "learned" in warns.lower() or "rule" in warns.lower() diff --git a/docs/train/core-inference.md b/docs/train/core-inference.md new file mode 100644 index 0000000..4e23376 --- /dev/null +++ b/docs/train/core-inference.md @@ -0,0 +1,92 @@ +# Core 加载 L1–L3 训练权重（#104） + +> 将 `train/` 导出的布局权重接入 `core` structure 管线；默认仍为几何 **rule**。 + +## 配置（环境变量，`ENPU_` 前缀） + +| 变量 | 默认 | 说明 | +|------|------|------| +| `ENPU_PIPELINE_MODE` | `legacy` | 设为 `structure` 才走 L1–L5 | +| `ENPU_STRUCTURE_L1L3_ENGINE` | `rule` | `rule` \| `learned` | +| `ENPU_L1L3_WEIGHTS` | _(空)_ | `layout_net.pt` / `best.pt` / `last.pt` 路径 | +| `ENPU_L1L3_DEVICE` | `cpu` | `cpu` 或 `cuda` | +| `ENPU_L1L3_FALLBACK` | `rule` | 加载/推理失败时：`rule` 或 `none`（抛错） | + +Settings 字段：`structure_l1l3_engine`、`l1l3_weights`、`l1l3_device`、`l1l3_fallback`。 + +## 权重格式 + +与训练导出一致： + +1. **`enpu_layout_net_v0`**（`train` `export_state_dict`） + - `format`, `model` (state_dict), `tasks`, `l2_heat_len`, `l3_heat_len` +2. **训练 ckpt**（`best.pt` / `last.pt`） + - `model` + `cfg`（含 `page_h/w`, `row_h/w`, heat 长度等） + +Core **不** import `train/` 包；网络结构在 `core/app/pipeline/structure/learned/model.py` 与训练侧对齐。 + +## 行为 + +```text +structure recognize + L1: 规则版面（title / key_time / score） + L2+L3: + engine=rule → 现有 OpenCV 谱行 + 分割线 + engine=learned → LayoutNet 热力 → systems + splits + → normalize_splits / splits_to_measures + L4–L5: 不变（几何 ROI + OCR） +``` + +- 坐标：全图像素，与 data-spec / 桌面一致。 +- `meta.parse_warnings` / `preprocess_steps` 含 `l1l3=learned|rule` 与 fallback 信息。 +- 权重缺失、torch 未装、推理异常 → 默认 **fallback rule**（可观察 warning）。 + +## 依赖策略 + +| 场景 | 依赖 | +|------|------| +| 默认 core / CI / 精简安装包 | **不需要** torch；`engine=rule` | +| 本机 learned 推理 | `pip install torch`（与训练环境可共用） | +| Windows 默认 NSIS 包 | **不**捆绑 torch；仅 rule | + +可选：在 `requirements.txt` 注释中说明；勿把 torch 写入 `requirements-ci.txt` / sidecar 默认。 + +## 使用示例 + +```powershell +cd D:\workspace\EnPu +$env:PYTHONPATH = ".\core" +$env:ENPU_PIPELINE_MODE = "structure" +$env:ENPU_RECOGNIZE_ENGINE = "mock" # 或 paddleocr +$env:ENPU_STRUCTURE_L1L3_ENGINE = "learned" +$env:ENPU_L1L3_WEIGHTS = ".\train\runs\mvp_l2_l3\best.pt" +$env:ENPU_L1L3_DEVICE = "cpu" + +# 启动 core 后 POST /v1/recognize +# 或对比脚本： +python scripts\eval_l1l3_engines.py --data samples\layout --weights train\runs\mvp_l2_l3\best.pt --out reports\l1l3_engines.json +``` + +## 模块路径 + +```text +core/app/pipeline/structure/learned/ + model.py # LayoutNet + loader.py # 权重加载与缓存 + postprocess.py # 热力 → 框 / split x + adapter.py # → PageLayout / splits 规范化 + infer_l1l3.py # 端到端 L2+L3 +pipeline.py # engine 分支 + fallback +``` + +## 限制 + +- MVP 模型仅 **L2+L3 热力**；L1 仍为规则（hybrid）。 +- 不承诺 learned 全面超过 rule；先接入与可评测。 +- 无 L4–L5 学习模型。 +- 训练 UI / Framework 与 core 仅通过 **权重文件** 交换。 + +## 相关 + +- #104 本能力 · #95 训练 Framework · #101 训练 UI · #93 data-spec · #94 模型方案 +- [l1-l3-model-design.md](./l1-l3-model-design.md) · [l1-l3-data-spec.md](./l1-l3-data-spec.md) diff --git a/scripts/eval_l1l3_engines.py b/scripts/eval_l1l3_engines.py new file mode 100644 index 0000000..2b29a2b --- /dev/null +++ b/scripts/eval_l1l3_engines.py @@ -0,0 +1,217 @@ +#!/usr/bin/env python3 +"""Compare structure L1–L3 rule vs learned engines on layout samples (#104). + +Example:: + + $env:PYTHONPATH = ".\\core" + $env:ENPU_L1L3_WEIGHTS = ".\\train\\runs\\mvp_l2_l3\\best.pt" + python scripts\\eval_l1l3_engines.py --data samples\\layout --out reports\\l1l3_engines.json +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path + +import numpy as np + +ROOT = Path(__file__).resolve().parents[1] +CORE = ROOT / "core" +if str(CORE) not in sys.path: + sys.path.insert(0, str(CORE)) + +from app.config import Settings, clear_settings_cache # noqa: E402 +from app.pipeline.structure.pipeline import run_structure_recognize # noqa: E402 + + +def _split_metrics(gt_xs: list[float], pred_xs: list[float], max_dist: float = 12.0) -> dict: + gt = sorted(float(x) for x in gt_xs) + pred = sorted(float(x) for x in pred_xs) + if not gt and not pred: + return { + "split_count_mae": 0.0, + "split_count_exact": 1.0, + "split_mean_abs_x_error": 0.0, + "n_gt": 0, + "n_pred": 0, + } + used: set[int] = set() + dists: list[float] = [] + tp = 0 + for gx in gt: + best_i, best_d = -1, max_dist + 1 + for i, px in enumerate(pred): + if i in used: + continue + d = abs(px - gx) + if d < best_d: + best_d, best_i = d, i + if best_i >= 0 and best_d <= max_dist: + used.add(best_i) + tp += 1 + dists.append(best_d) + return { + "split_count_mae": float(abs(len(pred) - len(gt))), + "split_count_exact": 1.0 if len(pred) == len(gt) else 0.0, + "split_mean_abs_x_error": float(sum(dists) / len(dists)) if dists else float("nan"), + "n_gt": len(gt), + "n_pred": len(pred), + "tp": tp, + "fp": len(pred) - tp, + "fn": len(gt) - tp, + } + + +def _gt_splits(layout: dict) -> list[float]: + xs: list[float] = [] + for row in (layout.get("l3") or {}).get("rows") or []: + for sp in row.get("splits") or []: + xs.append(float(sp["x"] if isinstance(sp, dict) else sp)) + return xs + + +def _pred_splits(structure) -> list[float]: + if structure is None: + return [] + bl = structure.barlines if hasattr(structure, "barlines") else structure.get("barlines") + out = [] + for b in bl or []: + if isinstance(b, dict) and b.get("x") is not None: + out.append(float(b["x"])) + return out + + +def _pred_n_systems(structure) -> int: + if structure is None: + return 0 + items = structure.items if hasattr(structure, "items") else structure.get("items") or [] + return sum(1 for it in items if (getattr(it, "layer", None) or it.get("layer")) == "L2") + + +def run_engine(data: bytes, engine: str, weights: str) -> dict: + clear_settings_cache() + os.environ["ENPU_PIPELINE_MODE"] = "structure" + os.environ["ENPU_RECOGNIZE_ENGINE"] = os.environ.get("ENPU_RECOGNIZE_ENGINE", "mock") + os.environ["ENPU_STRUCTURE_L1L3_ENGINE"] = engine + os.environ["ENPU_L1L3_FALLBACK"] = "rule" + if weights: + os.environ["ENPU_L1L3_WEIGHTS"] = weights + clear_settings_cache() + settings = Settings() + resp = run_structure_recognize(data, settings=settings, filename="eval.png") + return { + "ok": resp.ok, + "warnings": list(resp.meta.parse_warnings or [])[:30], + "n_systems": _pred_n_systems(resp.structure), + "splits": _pred_splits(resp.structure), + "engine_used": engine, + } + + +def main() -> int: + ap = argparse.ArgumentParser(description="rule vs learned L1–L3 eval (#104)") + ap.add_argument("--data", type=Path, default=ROOT / "samples" / "layout") + ap.add_argument( + "--weights", + type=Path, + default=ROOT / "train" / "runs" / "mvp_l2_l3" / "best.pt", + ) + ap.add_argument("--out", type=Path, default=ROOT / "reports" / "l1l3_engines.json") + ap.add_argument("--limit", type=int, default=20) + args = ap.parse_args() + + samples = sorted({p.parent for p in args.data.rglob("layout.json")})[: args.limit] + if not samples: + print("no layout samples under", args.data) + return 1 + + weights = str(args.weights) if args.weights.is_file() else "" + if not weights: + print("WARNING: weights missing; learned will fallback to rule:", args.weights) + + per_sample = [] + agg = { + "rule": {"count_mae": [], "exact": [], "x_err": []}, + "learned": {"count_mae": [], "exact": [], "x_err": []}, + } + + for sdir in samples: + layout = json.loads((sdir / "layout.json").read_text(encoding="utf-8")) + img_name = (layout.get("image") or {}).get("path") or "image.png" + img_path = sdir / img_name + if not img_path.is_file(): + continue + data = img_path.read_bytes() + gt_xs = _gt_splits(layout) + row = {"id": sdir.name, "n_gt_splits": len(gt_xs), "engines": {}} + for eng in ("rule", "learned"): + try: + pred = run_engine(data, eng, weights) + m = _split_metrics(gt_xs, pred["splits"]) + row["engines"][eng] = { + "n_systems": pred["n_systems"], + "n_pred_splits": len(pred["splits"]), + **m, + "warnings_head": (pred["warnings"] or [])[:5], + } + agg[eng]["count_mae"].append(m["split_count_mae"]) + agg[eng]["exact"].append(m["split_count_exact"]) + if m["split_mean_abs_x_error"] == m["split_mean_abs_x_error"]: + agg[eng]["x_err"].append(m["split_mean_abs_x_error"]) + except Exception as e: + row["engines"][eng] = {"error": str(e)} + per_sample.append(row) + print( + sdir.name, + "rule_mae=", + row["engines"].get("rule", {}).get("split_count_mae"), + "learned_mae=", + row["engines"].get("learned", {}).get("split_count_mae"), + ) + + def mean(xs: list[float]) -> float: + return float(sum(xs) / len(xs)) if xs else float("nan") + + summary = { + eng: { + "mean_split_count_mae": mean(agg[eng]["count_mae"]), + "mean_split_count_exact": mean(agg[eng]["exact"]), + "mean_abs_x_error": mean(agg[eng]["x_err"]), + "n": len(agg[eng]["count_mae"]), + } + for eng in ("rule", "learned") + } + report = { + "weights": weights, + "data": str(args.data), + "summary": summary, + "samples": per_sample, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + md = args.out.with_suffix(".md") + md.write_text( + "# rule vs learned L1–L3 (#104)\n\n" + f"- weights: `{weights}`\n" + f"- data: `{args.data}`\n\n" + "| engine | count_mae | exact | mean |x| err | n |\n" + "|--------|-----------|-------|----------------|---|\n" + + "\n".join( + f"| {e} | {summary[e]['mean_split_count_mae']:.3f} | " + f"{summary[e]['mean_split_count_exact']:.3f} | " + f"{summary[e]['mean_abs_x_error']:.3f} | {summary[e]['n']} |" + for e in ("rule", "learned") + ) + + "\n", + encoding="utf-8", + ) + print("wrote", args.out, "and", md) + print(json.dumps(summary, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())