diff --git a/README.md b/README.md index 4bca5d7..ba2a59a 100644 --- a/README.md +++ b/README.md @@ -58,11 +58,21 @@ L5 音符节点 音高数字 OCR（+几何兜底）+ 时值线 / 高低音点 | L3 | **主存分割线**，小节框由 L2 行边界 + 线推导；桌面拖线编辑 | | L5 | **音高数字**以 OCR 为主，并与同节点几何特征绑定 | | 开关 | `ENPU_PIPELINE_MODE=structure` 或 `legacy` | +| L1–L3 引擎 | 默认 **rule**（OpenCV）；可选 **learned**（#104，需本机 `torch` + 权重，**不进 CI / 精简安装包**） | 桌面在结构模式下可叠图查看 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 / 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`）。 +L1–L3 **布局训练与推理**（#92–#104）： + +| 项 | 链接 | +|----|------| +| 数据规范 | [docs/train/l1-l3-data-spec.md](./docs/train/l1-l3-data-spec.md) | +| 模型方案 | [docs/train/l1-l3-model-design.md](./docs/train/l1-l3-model-design.md) | +| Framework + 训练 UI | [`train/`](./train/)（`python scripts/run_ui.py`） | +| Core 加载权重 | [docs/train/core-inference.md](./docs/train/core-inference.md)（`ENPU_STRUCTURE_L1L3_ENGINE=learned`） | + +**CI**：`core/requirements-ci.txt` **不含 torch**；默认 `rule` 路径单测。learned 相关用例在无 torch 时自动 skip。 --- diff --git a/core/README.md b/core/README.md index 1d5236d..4dd71d1 100644 --- a/core/README.md +++ b/core/README.md @@ -96,11 +96,13 @@ 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_STRUCTURE_L1L3_ENGINE` | `rule` | `rule` \| `learned`（#104） | +| `ENPU_L1L3_WEIGHTS` | | `best.pt` / `layout_net.pt`（learned 时） | | `ENPU_L1L3_DEVICE` | `cpu` | `cpu` / `cuda` | | `ENPU_L1L3_FALLBACK` | `rule` | 失败时回退 rule 或 `none` | +**torch 可选**：CI（`requirements-ci.txt`）与默认安装包 **不**安装 torch。仅在本机需要 `learned` 时 `pip install torch`。详见 [core-inference.md](../docs/train/core-inference.md)。 + ## 流水线说明 ### legacy（默认，#3 / #10） diff --git a/core/app/pipeline/structure/learned/__init__.py b/core/app/pipeline/structure/learned/__init__.py index bbd0bd7..5c4e58f 100644 --- a/core/app/pipeline/structure/learned/__init__.py +++ b/core/app/pipeline/structure/learned/__init__.py @@ -1,12 +1,20 @@ """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). +Torch is **optional**. Importing this package must not require torch; +only ``run_learned_l2_l3`` / weight loading need it. """ -from app.pipeline.structure.learned.infer_l1l3 import ( - LearnedL1L3Error, - run_learned_l2_l3, -) +from __future__ import annotations __all__ = ["LearnedL1L3Error", "run_learned_l2_l3"] + + +def __getattr__(name: str): + if name in ("LearnedL1L3Error", "run_learned_l2_l3"): + from app.pipeline.structure.learned.infer_l1l3 import ( + LearnedL1L3Error, + run_learned_l2_l3, + ) + + return LearnedL1L3Error if name == "LearnedL1L3Error" else run_learned_l2_l3 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/core/app/pipeline/structure/learned/loader.py b/core/app/pipeline/structure/learned/loader.py index a25bebe..53680fb 100644 --- a/core/app/pipeline/structure/learned/loader.py +++ b/core/app/pipeline/structure/learned/loader.py @@ -7,7 +7,7 @@ from pathlib import Path from typing import Any -from app.pipeline.structure.learned.model import LayoutNet, LayoutNetConfig +from app.pipeline.structure.learned.model import LayoutNetConfig, build_layout_net logger = logging.getLogger(__name__) @@ -60,7 +60,7 @@ def load_layout_weights( path: str | Path, *, device: str = "cpu", -) -> tuple[LayoutNet, dict[str, Any]]: +) -> tuple[Any, dict[str, Any]]: """Load ``enpu_layout_net_v0`` or train ``best.pt`` / ``last.pt``. Returns (model.eval(), meta dict with format/path/cfg fields). @@ -82,7 +82,10 @@ def load_layout_weights( fmt = str(payload.get("format") or "train_ckpt") cfg = _cfg_from_payload(payload) - model = LayoutNet(cfg) + try: + model = build_layout_net(cfg) + except ImportError as e: + raise WeightsLoadError(str(e)) from e try: model.load_state_dict(payload["model"], strict=True) except Exception as e: @@ -111,7 +114,7 @@ def load_layout_weights( def get_cached_layout_model( path: str, device: str = "cpu", -) -> tuple[LayoutNet, dict[str, Any]]: +) -> tuple[Any, dict[str, Any]]: """Process-level cache keyed by path+device.""" return load_layout_weights(path, device=device) diff --git a/core/app/pipeline/structure/learned/model.py b/core/app/pipeline/structure/learned/model.py index 92a61e5..8381ec9 100644 --- a/core/app/pipeline/structure/learned/model.py +++ b/core/app/pipeline/structure/learned/model.py @@ -2,87 +2,15 @@ Architecture mirrors ``train/enpu_train/models/layout_net.py`` so exported weights load without importing the train package. + +Torch is imported lazily so CI (requirements-ci, no torch) can import sibling +modules (adapter/postprocess) without failing collection. """ 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) +from typing import Any @dataclass @@ -97,21 +25,114 @@ class LayoutNetConfig: 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 +def _torch_nn(): + try: + import torch + import torch.nn as nn + import torch.nn.functional as F + except ImportError as e: + raise ImportError( + "torch is required for learned L1–L3 weights. " + "Install with: pip install torch (optional; default engine is rule)" + ) from e + return torch, nn, F + + +def build_layout_net(cfg: LayoutNetConfig | None = None) -> Any: + """Construct LayoutNet (requires torch).""" + torch, nn, F = _torch_nn() + cfg = cfg or LayoutNetConfig() + + 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): # type: ignore[no-untyped-def] + 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): # type: ignore[no-untyped-def] + 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): # type: ignore[no-untyped-def] + 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) + + class LayoutNet(nn.Module): + def __init__(self, net_cfg: LayoutNetConfig) -> None: + super().__init__() + self.cfg = net_cfg + self.l2 = PageL2Head(net_cfg.l2_heat_len, net_cfg.base_channels) + self.l3 = RowL3Head(net_cfg.l3_heat_len, net_cfg.base_channels) + + def forward(self, page=None, rows=None): # type: ignore[no-untyped-def] + out = {} + 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 + + return LayoutNet(cfg) + + +# Back-compat name used by loader +class LayoutNet: # type: ignore[no-redef] + """Placeholder type name; construct via ``build_layout_net``.""" + + def __new__(cls, cfg: LayoutNetConfig | None = None): + return build_layout_net(cfg) diff --git a/core/tests/test_learned_l1l3.py b/core/tests/test_learned_l1l3.py index 7fb68d7..04e889b 100644 --- a/core/tests/test_learned_l1l3.py +++ b/core/tests/test_learned_l1l3.py @@ -1,4 +1,7 @@ -"""Tests for learned L1–L3 structure engine (#104).""" +"""Tests for learned L1–L3 structure engine (#104). + +Most tests run without torch (CI). Weight load / forward tests skip if torch missing. +""" from __future__ import annotations @@ -21,7 +24,15 @@ ) from app.pipeline.structure.pipeline import run_structure_recognize +try: + import torch as _torch_mod + + HAS_TORCH = True +except ImportError: + _torch_mod = None # type: ignore[assignment] + HAS_TORCH = False +requires_torch = pytest.mark.skipif(not HAS_TORCH, reason="torch not installed (optional)") def test_decode_peaks_and_boxes() -> None: heat = np.zeros(64, dtype=np.float32) heat[10] = 0.9 @@ -55,6 +66,7 @@ def test_load_missing_weights() -> None: load_layout_weights(Path("does_not_exist.pt")) +@requires_torch def test_load_real_weights_if_present() -> None: root = Path(__file__).resolve().parents[2] candidates = [ @@ -67,15 +79,11 @@ def test_load_real_weights_if_present() -> None: 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(): + page = _torch_mod.rand(1, 3, model.cfg.page_h, model.cfg.page_w) + with _torch_mod.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") @@ -123,6 +131,7 @@ def test_pipeline_learned_fallback_when_no_weights( assert any("fallback" in w.lower() or "failed" in w.lower() for w in warns) +@requires_torch def test_pipeline_learned_with_weights_if_present( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/docs/train/core-inference.md b/docs/train/core-inference.md index 4e23376..2e299f5 100644 --- a/docs/train/core-inference.md +++ b/docs/train/core-inference.md @@ -45,11 +45,13 @@ structure recognize | 场景 | 依赖 | |------|------| -| 默认 core / CI / 精简安装包 | **不需要** torch；`engine=rule` | +| 默认 core / **CI** / 精简安装包 | **不需要** torch；`engine=rule` | | 本机 learned 推理 | `pip install torch`（与训练环境可共用） | | Windows 默认 NSIS 包 | **不**捆绑 torch；仅 rule | -可选：在 `requirements.txt` 注释中说明；勿把 torch 写入 `requirements-ci.txt` / sidecar 默认。 +- `core/requirements-ci.txt` **不得**加入 torch。 +- `structure/learned` 在 **import 时不强制加载 torch**（懒加载）；仅 `load_layout_weights` / `run_learned_l2_l3` 需要。 +- 无 torch 时设 `ENPU_STRUCTURE_L1L3_ENGINE=learned` 会 **fallback 到 rule**（默认 `ENPU_L1L3_FALLBACK=rule`）。 ## 使用示例