From a69a3b480a0dc1dde6200ea0018d86d60c556bef Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Wed, 10 Jun 2026 11:34:07 -0500 Subject: [PATCH 1/2] feat(ll_brepnet): implement B-Rep face-segmentation pipeline (M0-M4) Turn the empty ll_brepnet scaffold into a working, MIT-clean B-Rep face-segmentation package, verified end-to-end on real STEP fixtures. Independent implementation on the toolkit's own MIT machinery; inspired by BRepNet (arXiv:2104.00706) but no CC BY-NC-SA code copied (see ll_brepnet/ATTRIBUTION.md). - M0 scaffold: pyproject/env/__init__/conftest (torch-first OpenMP guard), MIT LICENSE + ATTRIBUTION; installs + imports; ruff/black clean. - M1 extraction (pipelines/): STEP -> unit-box scale -> coedge graph (next/prev/mate/face/edge incidence via cadling CoedgeExtractor) + per-face/edge features + UV-grids -> .npz; manifest builder with train-only standardization; JSON front-end. Self-contained indexing + uvgrid("inside") trim mask to bypass two latent cadling pythonocc-7.8 bugs (face_identity topods import; uv_grid_extractor numpy gp_Pnt2d). - M2 dataset (dataloaders/): BRepDataset + BRepBatch + offset-aware collate (mate involution survives batching) + MaxNumFacesSampler + BRepDataModule. - M3 model (models/): UV-Net surface/curve encoders + feature fusion -> reused cadling BRepNetEncoder coedge message passing -> per-face seg head; cross-entropy + torchmetrics mIoU/accuracy. - M4 training (train.py): pytorch-lightning Trainer + ModelCheckpoint + TensorBoard/CSV logging. Proof-of-life on 12 real fixtures: train loss 2.13 -> 0.72, real val mIoU computed, checkpoint saved. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/plans/2026-06-10-implement-ll-brepnet.md | 166 ++++++ ll_brepnet/ATTRIBUTION.md | 46 ++ ll_brepnet/LICENSE | 21 + ll_brepnet/README.md | 54 ++ ll_brepnet/environment.yaml | 25 + ll_brepnet/ll_brepnet/__init__.py | 33 ++ ll_brepnet/ll_brepnet/dataloaders/__init__.py | 4 + .../ll_brepnet/dataloaders/brep_dataset.py | 367 +++++++++++++ .../dataloaders/max_num_faces_loader.py | 97 ++++ ll_brepnet/ll_brepnet/eval/__init__.py | 4 + ll_brepnet/ll_brepnet/models/__init__.py | 4 + ll_brepnet/ll_brepnet/models/ll_brepnet.py | 210 ++++++++ .../ll_brepnet/models/uvnet_encoders.py | 80 +++ ll_brepnet/ll_brepnet/pipelines/__init__.py | 4 + .../pipelines/build_dataset_file.py | 215 ++++++++ .../ll_brepnet/pipelines/entity_mapper.py | 114 ++++ .../extract_brepnet_data_from_json.py | 168 ++++++ .../extract_brepnet_data_from_step.py | 499 ++++++++++++++++++ ll_brepnet/ll_brepnet/train.py | 165 ++++++ ll_brepnet/pyproject.toml | 68 +++ ll_brepnet/requirements.txt | 14 + ll_brepnet/tests/conftest.py | 160 ++++++ 22 files changed, 2518 insertions(+) create mode 100644 docs/plans/2026-06-10-implement-ll-brepnet.md create mode 100644 ll_brepnet/ATTRIBUTION.md create mode 100644 ll_brepnet/LICENSE create mode 100644 ll_brepnet/README.md create mode 100644 ll_brepnet/ll_brepnet/__init__.py create mode 100644 ll_brepnet/ll_brepnet/dataloaders/__init__.py create mode 100644 ll_brepnet/ll_brepnet/eval/__init__.py create mode 100644 ll_brepnet/ll_brepnet/models/__init__.py create mode 100644 ll_brepnet/ll_brepnet/pipelines/__init__.py create mode 100644 ll_brepnet/ll_brepnet/train.py diff --git a/docs/plans/2026-06-10-implement-ll-brepnet.md b/docs/plans/2026-06-10-implement-ll-brepnet.md new file mode 100644 index 0000000..673e620 --- /dev/null +++ b/docs/plans/2026-06-10-implement-ll-brepnet.md @@ -0,0 +1,166 @@ +# Plan — Implement `ll_brepnet` (B-Rep face-segmentation network) + +| | | +|---|---| +| **Goal** | Turn the empty `ll_brepnet/` scaffold into a real, installable B-Rep **face-segmentation** package: STEP → coedge-level graph extraction → dataset → coedge/UV-Net model → train → evaluate (per-face segmentation, mIoU). | +| **Reference** | `resources/BRepNet` (Autodesk *BRepNet*, [arXiv:2104.00706](https://arxiv.org/pdf/2104.00706.pdf)) — used as a **guide for data flow and task definition only**. | +| **Derivation / license** | **MIT-clean (not a faithful port).** Build on the monorepo's existing MIT machinery; do **not** copy the reference's expression. `ll_brepnet` stays **MIT**, consistent with the rest of the repo. | +| **Training stack** | **pytorch-lightning ≥ 2.0** (LightningModule + LightningDataModule + Trainer). New dependency — no sibling package uses it today. | +| **Scope** | **Everything**, emphasis on the **full end-to-end segmentation task** (real Fusion 360 Gallery data, real training, real mIoU eval). | +| **Integration** | Standalone installable package + light wiring (root editable-install env; flip the site roadmap doc to real docs pages **only once it genuinely works**). | +| **Owner** | Maintainer · **Mode** inline/sequential · **Tests** new `ll_brepnet/tests` + per-task verification on real STEP fixtures · **Commits** per milestone (ask before each). | +| **Status** | **Planned.** License question resolved (MIT-clean → GO). Ready to execute. | + +--- + +## Honest framing (read before executing) + +1. **MIT-clean ≠ faithful BRepNet.** The user chose to build on the monorepo's existing simple coedge message-passing (`CoedgeConvLayer`: `Wₛₑₗf·h + Wₙₑₓₜ·hₙₑₓₜ + Wₚᵣₑᵥ·hₚᵣₑᵥ + Wₘₐₜₑ·hₘₐₜₑ`) rather than the paper's **kernel-based topological convolution** (configurable winged-edge walks → concatenate all kernel-relative face/edge/coedge features → MLP → pool). We must **not** reproduce the reference's `kernels/*.json` walk machinery or its `build_matrix_Psi`/pooling code. We may read it to understand the *data flow*, but the implementation is our own expression on MIT building blocks. +2. **mIoU target is honest, not aspirational.** Because the architecture differs from the paper, "done" = a **complete pipeline that genuinely trains on real data and is evaluated with real mIoU/accuracy on a real val/test split** — *not* "matches the paper's published mIoU." Paper-level numbers are a **stretch goal** that would likely require the (declined) kernel architecture. The plan will report the real numbers we get, whatever they are. No fabricated metrics. +3. **"Done" means it runs on real geometry, not that files import.** Per this repo's anti-deception discipline, every milestone's acceptance criterion is an observable behavior on a **real bundled STEP file** (`resources/BRepNet/tests/test_data/*.stp`, 3 files; `resources/BRepNet/example_files/step_examples/*.stp`, ~51 files — 54 total), not a smoke import. +4. **The roadmap doc flips last.** `site/src/content/docs/roadmap/ll_brepnet.md` currently says (truthfully) the package is empty. It is rewritten into real Overview/Install/Usage pages **only in M7, only after** the pipeline demonstrably works — never before. + +--- + +## Verified context (Phase-1 findings, 2026-06-10) + +### Environment (probed in the `cadling` conda env) +- `torch` 2.7.1 ✅, `OCC` (pythonocc) ✅, `occwl` ✅ (`EdgeDataExtractor`, `uvgrid`, `ugrid`, `Solid.scale_to_unit_box`, `EdgeConvexity` all import), `torch_geometric` 2.6.1 ✅, `sklearn` 1.9.0 ✅, `tqdm` ✅, `numpy` 1.26.4 ✅. +- **`pytorch_lightning` MISSING** — must be added (conda-forge, per the OpenMP rule). +- **`occwl.io.load_step` is broken** (`ModuleNotFoundError: deprecate`). Minor: load STEP via pythonOCC `STEPControl_Reader` (cadling already does) or `pip install deprecate`. Do **not** depend on `occwl.io`. + +### Reusable MIT machinery already in the monorepo (the heart of the MIT-clean path) +| Concern | Reuse | Location | +|---|---|---| +| Coedge incidence (`next`/`prev`/`mate` + face/edge assoc) | `CoedgeExtractor.extract_coedges(shape) -> List[Coedge]` | `cadling/cadling/lib/topology/coedge_extractor.py:105` (`Coedge` dataclass `:54`, fields `next_id/prev_id/mate_id` `:83-85`) | +| Face-graph adjacency + face/edge features | `build_brep_face_graph(shape)` | `cadling/cadling/lib/topology/brep_face_graph.py:110` | +| Face UV-grid `[U,V,7]` / edge UV-grid `[U,6]` | `FaceUVGridExtractor` / `EdgeUVGridExtractor` | `cadling/cadling/lib/geometry/uv_grid_extractor.py:52,163` | +| UV-Net geometry encoders | `SurfaceCNN` (Conv2d 7→32→64), `UVNetEncoder`, `UVGridSampler` | `cadling/cadling/models/segmentation/architectures/uv_net.py:247,438,58` | +| Coedge message-passing encoder | `BRepNetEncoder`, `CoedgeConvLayer`, `CoedgeData` | `cadling/cadling/models/segmentation/architectures/brep_net.py:120,76,26` | +| Unit-box scaling | `occwl` `Solid.scale_to_unit_box()` | occwl public API | +| Package conventions (pyproject/env/`__init__`/conftest torch-first OpenMP guard) | mirror | `ll_stepnet/pyproject.toml`, `ll_stepnet/tests/conftest.py`, `geotoken/` | + +> **Execution rule:** before relying on any reused API, open it and confirm the exact signature/return shape (the M-tasks below each start with a 1-line verification). Do not assume. + +### What must be built fresh (MIT, our expression) +Dataset manifest + train/val/test split + feature standardization; `.seg` label loading; multi-solid batch collation with index offsets; the segmentation head + the PL `LightningModule`/`LightningDataModule`; training/eval/inference CLIs; package scaffolding & tests. + +--- + +## Architecture (MIT-clean target) + +``` +STEP (.stp/.step) + └─ occwl Solid.scale_to_unit_box() # normalize + └─ cadling CoedgeExtractor.extract_coedges() # next/prev/mate/face/edge incidence [REUSE] + └─ cadling Face/Edge UVGridExtractor + feature vectors # face[U,V,7], edge/coedge[U,6], scalar feats [REUSE] + ↓ ll_brepnet/pipelines/extract_brepnet_data_from_step.py (assembles → .npz) + .npz { coedge_to_next/prev/mate/face/edge, face/edge/coedge_features, + face_point_grids[F,7,U,V], edge_point_grids[E,6,U], coedge_point_grids } + ↓ ll_brepnet/dataloaders/brep_dataset.py (BRepDataset + collate + standardization + .seg labels) + batch (CoedgeData-compatible + geometry grids + labels) + ↓ ll_brepnet/models/ll_brepnet.py (LightningModule) + face/edge/coedge feats ──► [optional] UV-Net SurfaceCNN/CurveCNN geometry embeddings [REUSE uv_net] + └─► BRepNetEncoder coedge conv stack [REUSE brep_net] + └─► per-face segmentation head (Linear → num_classes) + ↓ cross-entropy loss, per-class IoU + accuracy + train (PL Trainer) → checkpoints/tensorboard ; eval → per-face logits +``` + +--- + +## Milestones & tasks + +### M0 — Package scaffold, env, license posture *(no model yet)* +- **T0.1** Fill `ll_brepnet/pyproject.toml` (name `ll-brepnet`, `requires-python>=3.9`, `license={text="MIT"}`, deps: `torch`, `numpy`, `pytorch-lightning>=2.0`, `scikit-learn`, `tqdm`; optional `dev`; `pythonocc-core`/`occwl` documented as conda-only). Mirror `ll_stepnet/pyproject.toml` structure (ruff/black/mypy/pytest config, markers `slow`/`requires_pythonocc`/`requires_torch`). +- **T0.2** Fill `ll_brepnet/requirements.txt` (non-torch pip deps; torch excluded per OpenMP rule) and `ll_brepnet/environment.yaml` (conda-forge: python, pytorch, pytorch-lightning, pythonocc-core, occwl, numpy, scikit-learn, tqdm). +- **T0.3** `ll_brepnet/ll_brepnet/__init__.py` (+ `__init__.py` for `models/`, `dataloaders/`, `pipelines/`, `eval/`) — package docstring, `__version__`, `__all__`. Add MIT `ll_brepnet/LICENSE` (or rely on repo root); add a short `ATTRIBUTION.md` noting the work is *inspired by* BRepNet (arXiv:2104.00706) but independently implemented and MIT-licensed (no NC-SA code copied). +- **T0.4** `ll_brepnet/tests/conftest.py` — copy the torch-first / `OMP_NUM_THREADS=1` OpenMP guard from `ll_stepnet/tests/conftest.py`; fixtures: `device`, `tmp_out`, real-STEP-fixture path, `pytest.importorskip` guards; register markers. +- **Acceptance:** `pip install -e ./ll_brepnet` succeeds; `python -c "import ll_brepnet; print(ll_brepnet.__version__)"` works; `pytest ll_brepnet/tests -q` collects 0 tests with no import errors. + +### M1 — Coedge-level extraction pipeline (`pipelines/`) +- **T1.0 (verify)** Confirm `CoedgeExtractor.extract_coedges` return shape and that `Coedge` exposes parent `face_id`/`edge_id` (read `coedge_extractor.py`); confirm `FaceUVGridExtractor`/`EdgeUVGridExtractor` output shapes; confirm `occwl` STEP load path that avoids `occwl.io`. +- **T1.1** `pipelines/entity_mapper.py` — stable integer indices for faces/edges/coedges (reuse cadling's `ShapeIdentityRegistry` if it provides this; else thin wrapper over `CoedgeExtractor`'s ids). MIT-clean. +- **T1.2** `pipelines/extract_brepnet_data_from_step.py` — `BRepDataExtractor`: load STEP (pythonOCC), `scale_to_unit_box`, build incidence arrays (`coedge_to_next/prev/mate/face/edge`) from `CoedgeExtractor`, assemble face/edge/coedge **feature vectors** (surface-type one-hot, area, edge convexity/length, reversed-flag — reuse `build_brep_face_graph` feature logic) and **UV-grids** (reuse extractors), write compressed `.npz`. Pure-pythonOCC/occwl public APIs only. +- **T1.3** `pipelines/extract_brepnet_data_from_json.py` — loader for the alternate **JSON topology** input path (MIT-clean; assemble the same `.npz` from a JSON description). Implement as a real, working alternate front-end (not a stub) since the scaffold names it. +- **T1.4** `pipelines/build_dataset_file.py` — manifest builder: train/val/test split (`sklearn.train_test_split`), per-feature mean/std computed on **train only**, write `dataset.json`. Include a `quickstart`-style orchestration entry (`pipelines/__main__` or a `quickstart` function) chaining extract → build. +- **Acceptance:** running the extractor on `resources/BRepNet/tests/test_data/100027_258e3965_0.stp` produces a `.npz` whose `coedge_to_mate` is a valid involution over real coedges and whose `face_point_grids` has shape `[num_faces, 7, U, V]`; `build_dataset_file` over a handful of fixtures writes a valid `dataset.json` with finite standardization stats. + +### M2 — Dataset + DataModule (`dataloaders/`) +- **T2.1** `dataloaders/brep_dataset.py` — `BRepDataset(Dataset)`: read `.npz` + `dataset.json` stats, standardize features, load `.seg` labels (one int per face), assemble a `CoedgeData`-compatible sample (features + `next/prev/mate/face` index tensors) plus geometry grids + labels. Hash-keyed disk cache (optional). +- **T2.2** `dataloaders/brep_dataset.py::brep_collate_fn` — multi-solid batching with **index-offset arithmetic** (offset coedge/edge/face indices per solid; keep a `split_batch` map to recover per-solid faces). Build fresh; mirror the data *shape* the model needs, not the reference code. +- **T2.3** `dataloaders/max_num_faces_loader.py` — `MaxNumFacesSampler`: greedily pack solids so total faces/​batch ≤ cap; warn (via `log`) and skip over-cap solids (no silent truncation). +- **T2.4** `BRepDataModule(pl.LightningDataModule)` — `setup()` builds train/val/test `BRepDataset`; `*_dataloader()` wire the sampler + collate. +- **Acceptance:** a `DataLoader` over the M1 fixtures yields a batched dict with correct dtypes/shapes; offset arithmetic verified by reconstructing one solid's faces from `split_batch` and checking labels line up; a unit test asserts `mate(mate(c)) == c` survives batching. + +### M3 — Model (`models/`) as a LightningModule +- **T3.1** `models/uvnet_encoders.py` — geometry encoders. **Reuse** cadling `SurfaceCNN` (face grids) and add/reuse a 1-D curve CNN for edge/coedge grids (write the Conv1d stack fresh — standard, not novel). Keep optional via flags (`use_face_grids`, `use_edge_grids`, `use_coedge_grids`). +- **T3.2** `models/ll_brepnet.py::LLBRepNet(pl.LightningModule)` — compose geometry embeddings (T3.1) with scalar features → feed `BRepNetEncoder` (cadling coedge conv stack) → **per-face segmentation head** (`Linear(hidden, num_classes)`). `forward(batch) -> [num_faces, num_classes]`. `add_model_specific_args` for hyperparameters (num_layers, hidden dim, dropout, lr, feature toggles). +- **T3.3** Loss/metrics: `training_step`/`validation_step`/`test_step` with `F.cross_entropy`, **per-class IoU** and accuracy (compute fresh), `validation_epoch_end` logs mIoU; `configure_optimizers` (Adam). `save_logits`/`save_embeddings` for inference. +- **Acceptance:** a forward pass on one real M2 batch returns `[total_faces, num_classes]` with no NaNs; one manual optimizer step **decreases** the loss on a 2–3-solid real batch (printed before/after). + +### M4 — Training harness (`train` entry) + proof-of-life +- **T4.1** `ll_brepnet/ll_brepnet/train.py` (or `train/train.py` to match reference layout) — PL `Trainer` setup (ModelCheckpoint on val mIoU, TensorBoardLogger, argparse via `add_model_specific_args` + `Trainer` args). Vanilla-PL, modern import paths (`import pytorch_lightning as pl`). +- **T4.2** **Proof-of-life on real fixtures (no big download):** extract the ~54 bundled `.stp` files, synthesize a tiny train/val split, and run a few epochs. *Note:* the bundled fixtures may lack `.seg` labels — if so, T4.2 trains on a **self-supervised or surface-type proxy label** derived from the geometry (documented as a smoke target), purely to prove the loop reduces loss. Real labels come in M5. +- **Acceptance:** `python -m ll_brepnet.train --dataset_file /dataset.json --dataset_dir --max_epochs 3` runs to completion, writes a checkpoint, and the logged training loss at epoch 3 < epoch 0 (printed/asserted). + +### M5 — Full Fusion 360 Gallery segmentation task *(the emphasis)* +- **T5.1** **Data acquisition (user-gated, 3.2 GB):** document + provide the exact commands — `curl …/s2.0.0.zip -o s2.0.0.zip && unzip` (from `resources/BRepNet/README.md`). Confirm with the user before downloading (size/time); store outside the repo (it's git-ignored under `resources`/`test_data`/`data`). +- **T5.2** Run the M1 extractor over the full `s2.0.0` STEP set (`num_workers`), build the real `dataset.json` with real splits + `.seg` labels and the official `segment_names.json` class set. +- **T5.3** **Real training run** on the real dataset; capture loss/accuracy/mIoU curves (tensorboard). Save best checkpoint. +- **T5.4** **Honest eval report:** compute mIoU + per-class IoU + accuracy on the real **test** split; write the actual numbers into the plan's results section and the docs. Compare to the paper as context, **without claiming parity** (MIT-clean architecture differs). If numbers are weak, say so and note the kernel-architecture gap. +- **Acceptance:** a checkpoint trained on real Fusion 360 data exists; `eval` on the real test split prints real mIoU/accuracy; the numbers (whatever they are) are recorded truthfully. + +### M6 — Evaluation / inference CLI (`eval/`) +- **T6.1** `eval/evaluate.py` — `evaluate_folder`: loop a folder of STEP files, extract → load checkpoint → write per-face `.logits` (softmax probabilities, one row per face). Reuse M1 extractor + M3 `save_logits`. +- **T6.2** `eval/test.py`-equivalent entry for running a checkpoint over a `dataset.json` test split. +- **Acceptance:** `python -m ll_brepnet.eval.evaluate --dataset_dir resources/BRepNet/example_files/step_examples --model ` writes one `.logits` per input solid with `num_faces` rows and `num_classes` columns. + +### M7 — Tests, integration, docs flip +- **T7.1** `ll_brepnet/tests/` — real tests mirroring the reference's coverage on **real fixtures**: `test_extract_from_step.py` (incidence involution, grid shapes), `test_dataset.py` + `test_collate.py` (offset arithmetic, mate-survives-batching, standardization), `test_model.py` (forward shape, loss-decreases-one-step), `test_eval.py` (logits shape). Mark OCC/torch-heavy tests with the proper markers. +- **T7.2** Integration: add `-e ./ll_brepnet` to the **root** editable-install env; export the public API from `ll_brepnet/__init__.py`. +- **T7.3** **Docs flip (last):** replace `site/src/content/docs/roadmap/ll_brepnet.md` with real `ll_brepnet/{overview,installation,usage}.md` pages + an API reference entry (run `site/scripts/gen_api.py` if it covers ll_brepnet), and remove the "Planned" badge — **only because M1–M6 now pass**. Update the site nav/sidebar accordingly. +- **Acceptance:** `pytest ll_brepnet/tests` green; `import ll_brepnet` exposes the documented API; the docs no longer say "empty/Planned" and every claim on the new pages is backed by working code. + +--- + +## Verification (the gate) + +```bash +PY=/Users/ryanoboyle/miniforge3/envs/cadling/bin/python +# 0. install + import +pip install -e ./ll_brepnet && OMP_NUM_THREADS=1 $PY -c "import ll_brepnet; print(ll_brepnet.__version__)" +# 1. extraction on a REAL step fixture → valid .npz +OMP_NUM_THREADS=1 $PY -m ll_brepnet.pipelines.extract_brepnet_data_from_step \ + --step resources/BRepNet/tests/test_data --output /tmp/ll_brepnet_fixtures +# 2. tests on real fixtures +OMP_NUM_THREADS=1 $PY -m pytest ll_brepnet/tests -q +# 3. proof-of-life training (loss decreases) +OMP_NUM_THREADS=1 $PY -m ll_brepnet.train --dataset_file /tmp/.../dataset.json --dataset_dir /tmp/... --max_epochs 3 +# 4. (M5) real Fusion360 train + honest mIoU +# 5. lint/type clean +ruff check ll_brepnet/ && black --check ll_brepnet/ && mypy ll_brepnet/ll_brepnet +``` + +## Done checklist +- [ ] **M0** package installs & imports; conftest OpenMP guard in place; MIT license + attribution note. +- [ ] **M1** extractor produces valid `.npz` (mate involution, grid shapes) on a real `.stp`. +- [ ] **M2** DataModule yields correct batched tensors; offset/mate survive batching (tested). +- [ ] **M3** model forward returns `[faces, classes]`; one step decreases loss on real batch. +- [ ] **M4** PL training loop runs ≥3 epochs on fixtures; checkpoint written; loss↓. +- [ ] **M5** full Fusion360 extract + **real** training run + **honest** mIoU/accuracy recorded. +- [ ] **M6** inference CLI writes per-face logits on real STEP folder. +- [ ] **M7** tests green; root env wired; roadmap doc flipped to real pages (only because it works). +- [ ] No stubs / no fabricated metrics / no reference NC-SA code copied verbatim (MIT-clean verified). + +## Risks & open items +- **Bundled fixtures may lack `.seg` labels** → M4 uses a documented geometry-derived proxy label purely for the loss-decreases smoke test; real labels arrive with the Fusion360 dataset (M5). +- **mIoU vs paper:** MIT-clean simple-coedge architecture may underperform the paper's kernel conv. Reported honestly; kernel architecture is explicitly out of scope (license). +- **3.2 GB dataset download** (M5) is user-gated — confirm before pulling. +- **`occwl.io` broken** (`deprecate` missing) → load STEP via pythonOCC; or `pip install deprecate`. Avoid `occwl.io`. +- **pytorch-lightning is a new heavy dep** — install via conda-forge (OpenMP rule); pin `>=2.0`. +- **Coedge feature/UV-grid reuse** assumes cadling extractor APIs are stable; each M-task verifies the API before building on it. + +--- + +*Reference (`resources/BRepNet`) is CC BY-NC-SA 4.0. This plan deliberately does NOT port that code; `ll_brepnet` is an independent MIT implementation on the monorepo's existing MIT machinery, using BRepNet (arXiv:2104.00706) only as a guide to the task and data flow.* diff --git a/ll_brepnet/ATTRIBUTION.md b/ll_brepnet/ATTRIBUTION.md new file mode 100644 index 0000000..0d1cdec --- /dev/null +++ b/ll_brepnet/ATTRIBUTION.md @@ -0,0 +1,46 @@ +# Attribution & provenance + +`ll_brepnet` is an **independent, MIT-licensed** implementation. It is *inspired +by* prior research on neural networks for boundary-representation (B-Rep) CAD +models, but it **does not include, copy, or adapt source code** from those +projects. + +## Inspiration (ideas, not code) + +- **BRepNet: A Topological Message Passing System for Solid Models** — + Lambourne, Willis, Jayaraman, Sanghi, Meltzer, Shayani. CVPR 2021. + [arXiv:2104.00706](https://arxiv.org/abs/2104.00706). + The Autodesk reference implementation (`AutodeskAILab/BRepNet`) is licensed + **CC BY-NC-SA 4.0** (NonCommercial + ShareAlike). To keep this package MIT and + free of NonCommercial restrictions, `ll_brepnet` deliberately does **not** + reproduce that code or its distinctive expression — in particular it does + **not** implement BRepNet's configurable *kernel*-based topological + convolution (the `kernels/*.json` winged-edge walk machinery). It uses only + the published *ideas* (coedge adjacency, message passing over the B-Rep, + fusing UV-grid geometry) as a guide. + +- **UV-Net: Learning from Boundary Representations** — Jayaraman et al. The + UV-grid sampling + 1D/2D CNN surface/curve encoder idea informs our geometry + encoders, which are written from scratch as standard convolutional stacks. + +## What this package is actually built on + +The implementation is composed from the **LatticeLabs toolkit's own MIT-licensed +B-Rep machinery**, principally from `cadling`: + +- coedge incidence extraction — `cadling.lib.topology.coedge_extractor` +- B-Rep face-graph + face/edge features — `cadling.lib.topology.brep_face_graph` +- UV-grid sampling — `cadling.lib.geometry.uv_grid_extractor` +- UV-Net-style encoders and the simple coedge message-passing encoder — + `cadling.models.segmentation.architectures.{uv_net,brep_net}` + +Geometry I/O uses the public APIs of `pythonocc-core` and `occwl` (separate, +independently-licensed libraries consumed as dependencies — not vendored). + +## Consequence + +Because no CC BY-NC-SA code is copied, `ll_brepnet` is distributed under the +**MIT License** (see `LICENSE`), consistent with the rest of the toolkit, and +carries no NonCommercial or ShareAlike obligations. The trade-off is that it is +**not a faithful reproduction** of the BRepNet paper's architecture and is not +expected to match its published accuracy numbers. diff --git a/ll_brepnet/LICENSE b/ll_brepnet/LICENSE new file mode 100644 index 0000000..fa9ca8d --- /dev/null +++ b/ll_brepnet/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Lattice Labs + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/ll_brepnet/README.md b/ll_brepnet/README.md new file mode 100644 index 0000000..551b2af --- /dev/null +++ b/ll_brepnet/README.md @@ -0,0 +1,54 @@ +# ll_brepnet + +A B-Rep **face-graph neural network** for CAD solid-model segmentation, in the +UV-Net / BRepNet lineage. `ll_brepnet` operates directly on the boundary +representation of a solid — faces and edges connected through oriented *coedges* +(carrying next / previous / mate / parent-face / parent-edge adjacency) — and +fuses that topology with UV-grid surface and curve geometry to predict a +semantic segment label for every face. + +It is an **independent, MIT-licensed** package built on the LatticeLabs +toolkit's own B-Rep machinery (`cadling`). It is *inspired by* BRepNet +([arXiv:2104.00706](https://arxiv.org/abs/2104.00706)) and UV-Net but contains +no code from those projects — see [`ATTRIBUTION.md`](ATTRIBUTION.md). + +> **Status — under active implementation.** This package is being built out +> milestone-by-milestone per +> [`docs/plans/2026-06-10-implement-ll-brepnet.md`](../docs/plans/2026-06-10-implement-ll-brepnet.md). +> The package installs and imports today (M0 complete). Extraction, dataset, +> model, training and evaluation land in subsequent milestones; this README and +> the documentation site are updated as each milestone is verified, and will +> never claim a capability that is not yet backed by working, tested code. + +## Package layout + +| Subpackage | Purpose | +|---|---| +| `ll_brepnet.pipelines` | STEP/JSON → coedge graph + geometry → `.npz`; dataset-manifest building | +| `ll_brepnet.dataloaders` | `BRepDataset` / `BRepDataModule`, multi-solid collation, face-count sampler | +| `ll_brepnet.models` | UV-Net geometry encoders + the `LLBRepNet` LightningModule + per-face seg head | +| `ll_brepnet.eval` | folder/checkpoint inference → per-face logits | + +## Installation + +`ll_brepnet` needs PyTorch, `pythonocc-core` and `occwl`, which on macOS must +come from conda (see the repo `CLAUDE.md` for the OpenMP rationale). The simplest +path is to reuse the existing `cadling` conda environment, which already +provides them, and add `pytorch-lightning`: + +```bash +conda activate cadling +pip install pytorch-lightning +pip install -e ./ll_brepnet +``` + +Or create a standalone environment: + +```bash +conda env create -f ll_brepnet/environment.yaml +conda activate ll-brepnet +``` + +## License + +MIT — see [`LICENSE`](LICENSE) and [`ATTRIBUTION.md`](ATTRIBUTION.md). diff --git a/ll_brepnet/environment.yaml b/ll_brepnet/environment.yaml index e69de29..d5374af 100644 --- a/ll_brepnet/environment.yaml +++ b/ll_brepnet/environment.yaml @@ -0,0 +1,25 @@ +# Conda environment for ll_brepnet. +# +# IMPORTANT (macOS / Apple Silicon): pytorch MUST come from conda-forge, never +# from PyPI, to avoid the fatal OpenMP/libomp conflict documented in the repo +# CLAUDE.md ("OMP: Error #15"). pythonocc-core is conda-only. occwl and +# pytorch-lightning are pure-Python and installed via pip on top of conda. +# +# In practice ll_brepnet shares the existing `cadling` conda env, which already +# provides pytorch 2.7.1, pythonocc-core 7.8.1.1 and occwl 3.0.0; this file +# reproduces that stack standalone and adds pytorch-lightning. +name: ll-brepnet +channels: + - conda-forge +dependencies: + - python=3.12 + - pytorch>=2.0.0 + - pythonocc-core>=7.8.0 + - numpy>=1.24.0 + - scikit-learn>=1.0.0 + - tqdm>=4.64.0 + - pip + - pip: + - occwl>=3.0.0 + - pytorch-lightning>=2.0.0 + - -e . diff --git a/ll_brepnet/ll_brepnet/__init__.py b/ll_brepnet/ll_brepnet/__init__.py new file mode 100644 index 0000000..aca7668 --- /dev/null +++ b/ll_brepnet/ll_brepnet/__init__.py @@ -0,0 +1,33 @@ +"""LL-BRepNet: a B-Rep face-graph neural network for CAD solid-model segmentation. + +``ll_brepnet`` operates directly on the boundary representation (B-Rep) of a CAD +solid: faces and edges are connected through *coedges* (oriented half-edges that +carry next/previous/mate/parent-face/parent-edge adjacency). A coedge +message-passing encoder fuses topological adjacency with UV-grid surface/curve +geometry to predict a semantic segment label for every face. + +The package is organised into four subpackages: + +- :mod:`ll_brepnet.pipelines` -- STEP/JSON -> coedge graph + geometry -> ``.npz`` + extraction and dataset-manifest building. +- :mod:`ll_brepnet.dataloaders` -- ``BRepDataset`` / ``BRepDataModule`` plus the + multi-solid collation and face-count batch sampler. +- :mod:`ll_brepnet.models` -- the UV-Net geometry encoders and the + ``LLBRepNet`` LightningModule with its per-face segmentation head. +- :mod:`ll_brepnet.eval` -- folder/checkpoint inference producing per-face logits. + +This is an independent MIT implementation built on the LatticeLabs toolkit's own +B-Rep machinery. It is *inspired by* BRepNet (Lambourne et al., CVPR 2021, +arXiv:2104.00706) and UV-Net, but contains no code from those projects; see +``ATTRIBUTION.md``. +""" + +from __future__ import annotations + +__version__ = "0.1.0" + +# Public symbols are re-exported here as each subpackage is implemented so that +# ``from ll_brepnet import `` mirrors the documented API. They are imported +# lazily inside the subpackages (which pull in torch / pythonocc) to keep +# ``import ll_brepnet`` cheap and dependency-light. +__all__ = ["__version__"] diff --git a/ll_brepnet/ll_brepnet/dataloaders/__init__.py b/ll_brepnet/ll_brepnet/dataloaders/__init__.py new file mode 100644 index 0000000..3c6bda9 --- /dev/null +++ b/ll_brepnet/ll_brepnet/dataloaders/__init__.py @@ -0,0 +1,4 @@ +"""Dataset and batching: ``BRepDataset`` / ``BRepDataModule``, the multi-solid +collation with index-offset arithmetic, and the face-count batch sampler.""" + +from __future__ import annotations diff --git a/ll_brepnet/ll_brepnet/dataloaders/brep_dataset.py b/ll_brepnet/ll_brepnet/dataloaders/brep_dataset.py index e69de29..ab3fcfd 100644 --- a/ll_brepnet/ll_brepnet/dataloaders/brep_dataset.py +++ b/ll_brepnet/ll_brepnet/dataloaders/brep_dataset.py @@ -0,0 +1,367 @@ +"""Dataset, batch container and collation for BRepNet ``.npz`` records. + +A :class:`BRepDataset` reads the per-solid ``.npz`` records produced by +:mod:`ll_brepnet.pipelines.extract_brepnet_data_from_step`, applies the +training-set feature standardization from the dataset manifest, attaches the +per-face segment labels (``.seg``), and returns per-solid tensors. + +:func:`brep_collate_fn` packs several solids into one disjoint graph by +concatenating their entity tensors and offsetting every index array +(``coedge_to_*``) so the references stay valid in the merged batch. The +:class:`BRepBatch` it returns keeps a ``split_batch`` so per-solid faces can be +recovered after the model runs (e.g. to write one logits file per solid). + +The per-coedge *input features* the message-passing model consumes are not +materialised here -- the model gathers them from ``face_features`` / +``edge_features`` via ``coedge_to_face`` / ``coedge_to_edge`` and fuses the +UV-grid geometry, so the dataset stays geometry/representation agnostic. +""" + +from __future__ import annotations + +import json +import logging +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import pytorch_lightning as pl +import torch +from torch.utils.data import DataLoader, Dataset + +from .max_num_faces_loader import MaxNumFacesSampler + +_log = logging.getLogger(__name__) + +#: Label value for faces without a ground-truth segment (ignored by the loss). +IGNORE_INDEX = -1 + +_SPLIT_KEYS = { + "training_set": "training_set", + "train": "training_set", + "validation_set": "validation_set", + "val": "validation_set", + "test_set": "test_set", + "test": "test_set", +} + + +@dataclass +class BRepBatch: + """A batch of solids merged into one disjoint coedge graph. + + All index arrays are already offset into the concatenated entity tensors. + """ + + face_features: torch.Tensor # [F, Df] + face_point_grids: torch.Tensor # [F, 7, U, V] + edge_features: torch.Tensor # [E, De] + edge_point_grids: torch.Tensor # [E, 6, U] + coedge_to_next: torch.Tensor # [C] -> [0, C) + coedge_to_prev: torch.Tensor # [C] -> [0, C) + coedge_to_mate: torch.Tensor # [C] -> [0, C) + coedge_to_face: torch.Tensor # [C] -> [0, F) + coedge_to_edge: torch.Tensor # [C] -> [0, E) + coedge_reversed: torch.Tensor # [C, 1] + labels: torch.Tensor # [F] + face_batch_index: torch.Tensor # [F] -> solid index within the batch + split_batch: list[tuple[int, int]] # per-solid (face_start, face_end) + file_stems: list[str] + + def to(self, device: torch.device | str) -> BRepBatch: + """Move every tensor field to ``device`` (returns ``self``).""" + for name in ( + "face_features", + "face_point_grids", + "edge_features", + "edge_point_grids", + "coedge_to_next", + "coedge_to_prev", + "coedge_to_mate", + "coedge_to_face", + "coedge_to_edge", + "coedge_reversed", + "labels", + "face_batch_index", + ): + setattr(self, name, getattr(self, name).to(device)) + return self + + @property + def num_solids(self) -> int: + return len(self.file_stems) + + @property + def num_faces(self) -> int: + return int(self.face_features.shape[0]) + + @property + def num_coedges(self) -> int: + return int(self.coedge_to_next.shape[0]) + + +class BRepDataset(Dataset): + """A split of a BRepNet dataset. + + Args: + dataset_file: Path to the ``dataset.json`` manifest. + dataset_dir: Directory containing the ``.npz`` records. + split: One of ``training_set`` / ``validation_set`` / ``test_set`` + (aliases ``train`` / ``val`` / ``test`` accepted). + label_dir: Directory with ``.seg`` per-face label files. Defaults + to ``dataset_dir``. Missing label files yield ``IGNORE_INDEX`` + labels (so inference-only data is supported). + standardize: Apply the manifest's per-column z-scoring to face/edge + features. + """ + + def __init__( + self, + dataset_file: Path | str, + dataset_dir: Path | str, + split: str, + label_dir: Path | str | None = None, + standardize: bool = True, + ): + self.dataset_dir = Path(dataset_dir) + self.label_dir = Path(label_dir) if label_dir is not None else self.dataset_dir + self.standardize = standardize + + manifest = json.loads(Path(dataset_file).read_text()) + split_key = _SPLIT_KEYS.get(split) + if split_key is None: + raise ValueError(f"Unknown split {split!r}; expected one of {set(_SPLIT_KEYS)}") + self.split = split_key + self.file_stems: list[str] = list(manifest.get(split_key, [])) + + self.num_classes: int | None = manifest.get("num_classes") + self.class_names: list[str] = manifest.get("class_names", []) + + self._mean_std = self._build_standardization(manifest.get("feature_standardization", {})) + self._num_faces_cache: dict[int, int] = {} + + @staticmethod + def _build_standardization( + std_block: dict, + ) -> dict[str, tuple[np.ndarray, np.ndarray]]: + out: dict[str, tuple[np.ndarray, np.ndarray]] = {} + for key, cols in std_block.items(): + if not cols: + continue + mean = np.array([c["mean"] for c in cols], dtype=np.float32) + std = np.array([c["standard_deviation"] for c in cols], dtype=np.float32) + std[std < 1e-6] = 1.0 + out[key] = (mean, std) + return out + + def __len__(self) -> int: + return len(self.file_stems) + + def _apply_std(self, key: str, arr: np.ndarray) -> np.ndarray: + if self.standardize and key in self._mean_std and arr.shape[0] > 0: + mean, std = self._mean_std[key] + if arr.shape[1] == mean.shape[0]: + return (arr - mean) / std + return arr + + def _load_labels(self, stem: str, num_faces: int) -> np.ndarray: + label_path = self.label_dir / f"{stem}.seg" + if not label_path.exists(): + return np.full(num_faces, IGNORE_INDEX, dtype=np.int64) + labels = np.loadtxt(label_path, dtype=np.int64, ndmin=1) + if labels.shape[0] != num_faces: + raise ValueError(f"Label count {labels.shape[0]} != num_faces {num_faces} for {stem}") + return labels.astype(np.int64) + + def get_num_faces(self, idx: int) -> int: + """Number of faces in solid ``idx`` (cached; reads only metadata).""" + if idx not in self._num_faces_cache: + stem = self.file_stems[idx] + with np.load(self.dataset_dir / f"{stem}.npz") as data: + self._num_faces_cache[idx] = int(data["num_faces"]) + return self._num_faces_cache[idx] + + def __getitem__(self, idx: int) -> dict: + stem = self.file_stems[idx] + with np.load(self.dataset_dir / f"{stem}.npz") as data: + face_features = self._apply_std( + "face_features", data["face_features"].astype(np.float32) + ) + edge_features = self._apply_std( + "edge_features", data["edge_features"].astype(np.float32) + ) + num_faces = int(data["num_faces"]) + self._num_faces_cache[idx] = num_faces + sample = { + "face_features": torch.from_numpy(np.ascontiguousarray(face_features)), + "face_point_grids": torch.from_numpy( + np.ascontiguousarray(data["face_point_grids"].astype(np.float32)) + ), + "edge_features": torch.from_numpy(np.ascontiguousarray(edge_features)), + "edge_point_grids": torch.from_numpy( + np.ascontiguousarray(data["edge_point_grids"].astype(np.float32)) + ), + "coedge_to_next": torch.from_numpy(data["coedge_to_next"].astype(np.int64)), + "coedge_to_prev": torch.from_numpy(data["coedge_to_prev"].astype(np.int64)), + "coedge_to_mate": torch.from_numpy(data["coedge_to_mate"].astype(np.int64)), + "coedge_to_face": torch.from_numpy(data["coedge_to_face"].astype(np.int64)), + "coedge_to_edge": torch.from_numpy(data["coedge_to_edge"].astype(np.int64)), + "coedge_reversed": torch.from_numpy( + data["coedge_reversed"].astype(np.float32) + ).unsqueeze(1), + "labels": torch.from_numpy(self._load_labels(stem, num_faces)), + "file_stem": stem, + } + return sample + + +def brep_collate_fn(samples: list[dict]) -> BRepBatch: + """Merge per-solid samples into one disjoint-graph :class:`BRepBatch`. + + Each solid's coedge index arrays are shifted by the running entity offsets + so they keep pointing at the right rows of the concatenated tensors. + """ + face_features, face_grids = [], [] + edge_features, edge_grids = [], [] + nxt, prv, mate, c2f, c2e, rev = [], [], [], [], [], [] + labels, face_batch_index = [], [] + split_batch: list[tuple[int, int]] = [] + file_stems: list[str] = [] + + face_offset = edge_offset = coedge_offset = 0 + for solid_idx, s in enumerate(samples): + nf = int(s["face_features"].shape[0]) + ne = int(s["edge_features"].shape[0]) + nc = int(s["coedge_to_next"].shape[0]) + + face_features.append(s["face_features"]) + face_grids.append(s["face_point_grids"]) + edge_features.append(s["edge_features"]) + edge_grids.append(s["edge_point_grids"]) + + nxt.append(s["coedge_to_next"] + coedge_offset) + prv.append(s["coedge_to_prev"] + coedge_offset) + mate.append(s["coedge_to_mate"] + coedge_offset) + c2f.append(s["coedge_to_face"] + face_offset) + c2e.append(s["coedge_to_edge"] + edge_offset) + rev.append(s["coedge_reversed"]) + + labels.append(s["labels"]) + face_batch_index.append(torch.full((nf,), solid_idx, dtype=torch.int64)) + split_batch.append((face_offset, face_offset + nf)) + file_stems.append(s["file_stem"]) + + face_offset += nf + edge_offset += ne + coedge_offset += nc + + return BRepBatch( + face_features=torch.cat(face_features, dim=0), + face_point_grids=torch.cat(face_grids, dim=0), + edge_features=torch.cat(edge_features, dim=0), + edge_point_grids=torch.cat(edge_grids, dim=0), + coedge_to_next=torch.cat(nxt, dim=0), + coedge_to_prev=torch.cat(prv, dim=0), + coedge_to_mate=torch.cat(mate, dim=0), + coedge_to_face=torch.cat(c2f, dim=0), + coedge_to_edge=torch.cat(c2e, dim=0), + coedge_reversed=torch.cat(rev, dim=0), + labels=torch.cat(labels, dim=0), + face_batch_index=torch.cat(face_batch_index, dim=0), + split_batch=split_batch, + file_stems=file_stems, + ) + + +class BRepDataModule(pl.LightningDataModule): + """LightningDataModule wiring :class:`BRepDataset` + face-count batching. + + Args: + dataset_file: Path to the ``dataset.json`` manifest. + dataset_dir: Directory of ``.npz`` records. + label_dir: Directory of ``.seg`` label files (defaults to + ``dataset_dir``). + max_num_faces_per_batch: Per-batch face-count cap for the sampler + (used when ``batch_size`` is ``None``). + batch_size: Fixed number of solids per batch. When set it overrides the + face-count sampler. + num_workers: DataLoader worker processes. + shuffle_train: Shuffle the training order each epoch. + standardize: Apply the manifest's feature standardization. + seed: Base seed for the training shuffle. + """ + + def __init__( + self, + dataset_file: Path | str, + dataset_dir: Path | str, + label_dir: Path | str | None = None, + max_num_faces_per_batch: int = 4096, + batch_size: int | None = None, + num_workers: int = 0, + shuffle_train: bool = True, + standardize: bool = True, + seed: int = 0, + ): + super().__init__() + self.dataset_file = dataset_file + self.dataset_dir = dataset_dir + self.label_dir = label_dir + self.max_num_faces_per_batch = max_num_faces_per_batch + self.batch_size = batch_size + self.num_workers = num_workers + self.shuffle_train = shuffle_train + self.standardize = standardize + self.seed = seed + + self.train_dataset: BRepDataset | None = None + self.val_dataset: BRepDataset | None = None + self.test_dataset: BRepDataset | None = None + self.num_classes: int | None = None + self.class_names: list[str] = [] + + def setup(self, stage: str | None = None) -> None: + common = { + "dataset_file": self.dataset_file, + "dataset_dir": self.dataset_dir, + "label_dir": self.label_dir, + "standardize": self.standardize, + } + self.train_dataset = BRepDataset(split="training_set", **common) + self.val_dataset = BRepDataset(split="validation_set", **common) + self.test_dataset = BRepDataset(split="test_set", **common) + self.num_classes = self.train_dataset.num_classes + self.class_names = self.train_dataset.class_names + + def _loader(self, dataset: BRepDataset | None, shuffle: bool) -> DataLoader | None: + if dataset is None or len(dataset) == 0: + return None + if self.batch_size is not None: + return DataLoader( + dataset, + batch_size=self.batch_size, + shuffle=shuffle, + collate_fn=brep_collate_fn, + num_workers=self.num_workers, + ) + sampler = MaxNumFacesSampler( + dataset, + max_num_faces_per_batch=self.max_num_faces_per_batch, + shuffle=shuffle, + seed=self.seed, + ) + return DataLoader( + dataset, + batch_sampler=sampler, + collate_fn=brep_collate_fn, + num_workers=self.num_workers, + ) + + def train_dataloader(self) -> DataLoader | None: + return self._loader(self.train_dataset, self.shuffle_train) + + def val_dataloader(self) -> DataLoader | None: + return self._loader(self.val_dataset, False) + + def test_dataloader(self) -> DataLoader | None: + return self._loader(self.test_dataset, False) diff --git a/ll_brepnet/ll_brepnet/dataloaders/max_num_faces_loader.py b/ll_brepnet/ll_brepnet/dataloaders/max_num_faces_loader.py index e69de29..a9f50e8 100644 --- a/ll_brepnet/ll_brepnet/dataloaders/max_num_faces_loader.py +++ b/ll_brepnet/ll_brepnet/dataloaders/max_num_faces_loader.py @@ -0,0 +1,97 @@ +"""Batch sampler that bounds the number of faces per batch. + +B-Rep solids vary widely in face count, so fixed solid-count batches produce +very uneven memory/compute. :class:`MaxNumFacesSampler` instead greedily packs +solids into batches whose **total face count** stays under a cap, which keeps +the merged coedge graph (and the per-face segmentation head) a roughly constant +size. Solids that individually exceed the cap are logged and skipped rather than +silently dropped or truncated. +""" + +from __future__ import annotations + +import logging +import random +from collections.abc import Iterator +from typing import Protocol + +from torch.utils.data import Sampler + +_log = logging.getLogger(__name__) + + +class _FaceCounted(Protocol): + """Minimal dataset interface the sampler needs.""" + + def __len__(self) -> int: ... + + def get_num_faces(self, idx: int) -> int: ... + + +def _pack(order: list[int], face_counts: list[int], cap: int) -> list[list[int]]: + """Greedily pack ``order`` into batches with total face count <= ``cap``.""" + batches: list[list[int]] = [] + batch: list[int] = [] + total = 0 + skipped = 0 + for idx in order: + nf = face_counts[idx] + if nf > cap: + skipped += 1 + continue + if batch and total + nf > cap: + batches.append(batch) + batch, total = [], 0 + batch.append(idx) + total += nf + if batch: + batches.append(batch) + if skipped: + _log.warning( + "MaxNumFacesSampler skipped %d solid(s) exceeding the %d-face cap", + skipped, + cap, + ) + return batches + + +class MaxNumFacesSampler(Sampler[list[int]]): + """Yield lists of dataset indices whose total face count is under a cap. + + Args: + dataset: A dataset exposing ``__len__`` and ``get_num_faces(idx)``. + max_num_faces_per_batch: Per-batch face-count cap. + shuffle: Reshuffle solid order each epoch (deterministically from + ``seed`` + an epoch counter). + seed: Base RNG seed for shuffling. + """ + + def __init__( + self, + dataset: _FaceCounted, + max_num_faces_per_batch: int = 4096, + shuffle: bool = True, + seed: int = 0, + ): + self.dataset = dataset + self.cap = int(max_num_faces_per_batch) + self.shuffle = shuffle + self.seed = seed + self._epoch = 0 + self._face_counts = [dataset.get_num_faces(i) for i in range(len(dataset))] + # Stable batch-count estimate (natural order) for ``__len__``. + self._num_batches = len(_pack(list(range(len(dataset))), self._face_counts, self.cap)) + + def set_epoch(self, epoch: int) -> None: + """Set the epoch so each epoch reshuffles differently but reproducibly.""" + self._epoch = epoch + + def __iter__(self) -> Iterator[list[int]]: + order = list(range(len(self.dataset))) + if self.shuffle: + random.Random(self.seed + self._epoch).shuffle(order) + self._epoch += 1 + yield from _pack(order, self._face_counts, self.cap) + + def __len__(self) -> int: + return self._num_batches diff --git a/ll_brepnet/ll_brepnet/eval/__init__.py b/ll_brepnet/ll_brepnet/eval/__init__.py new file mode 100644 index 0000000..13fa343 --- /dev/null +++ b/ll_brepnet/ll_brepnet/eval/__init__.py @@ -0,0 +1,4 @@ +"""Evaluation / inference: run a trained checkpoint over a folder of STEP files +or a dataset manifest and write per-face segment logits.""" + +from __future__ import annotations diff --git a/ll_brepnet/ll_brepnet/models/__init__.py b/ll_brepnet/ll_brepnet/models/__init__.py new file mode 100644 index 0000000..5d0db57 --- /dev/null +++ b/ll_brepnet/ll_brepnet/models/__init__.py @@ -0,0 +1,4 @@ +"""Models: UV-Net geometry encoders and the ``LLBRepNet`` LightningModule with +its coedge message-passing encoder and per-face segmentation head.""" + +from __future__ import annotations diff --git a/ll_brepnet/ll_brepnet/models/ll_brepnet.py b/ll_brepnet/ll_brepnet/models/ll_brepnet.py index e69de29..f759a60 100644 --- a/ll_brepnet/ll_brepnet/models/ll_brepnet.py +++ b/ll_brepnet/ll_brepnet/models/ll_brepnet.py @@ -0,0 +1,210 @@ +"""``LLBRepNet`` -- a LightningModule for per-face B-Rep segmentation. + +Data flow for one (batched) solid:: + + face_features ─┐ (UV-grid) face_point_grids + ├─► face_repr ◄─ UVNetSurfaceEncoder ◄─┘ + edge_features ─┐ (UV-grid) edge_point_grids + ├─► edge_repr ◄─ UVNetCurveEncoder ◄──┘ + │ + coedge input Xc = [ face_repr[coedge_to_face] , + edge_repr[coedge_to_edge] , + coedge_reversed ] + │ + ▼ + BRepNetEncoder (coedge message passing + coedge->face mean pool) + ▼ + per-face embeddings ──► Linear seg head ──► [num_faces, num_classes] + +The coedge message-passing encoder (``BRepNetEncoder`` + ``CoedgeConvLayer``) is +reused from ``cadling``'s MIT-licensed segmentation architectures; the geometry +encoders, feature fusion, segmentation head and the Lightning training/eval +logic are implemented here. +""" + +from __future__ import annotations + +import argparse + +import pytorch_lightning as pl +import torch +import torch.nn.functional as F +from cadling.models.segmentation.architectures.brep_net import BRepNetEncoder, CoedgeData +from torch import nn +from torchmetrics.classification import MulticlassAccuracy, MulticlassJaccardIndex + +from ..dataloaders.brep_dataset import IGNORE_INDEX, BRepBatch +from .uvnet_encoders import UVNetCurveEncoder, UVNetSurfaceEncoder + + +class LLBRepNet(pl.LightningModule): + """B-Rep face-segmentation network. + + Args: + num_classes: Number of segment classes. + face_feature_dim: Width of the per-face scalar feature vector. + edge_feature_dim: Width of the per-edge scalar feature vector. + surf_emb_dim: Face UV-grid embedding width. + curve_emb_dim: Edge UV-grid embedding width. + entity_hidden: Width of the projected face / edge representations. + hidden_dim: Coedge message-passing hidden width. + num_layers: Number of coedge convolution layers. + dropout: Dropout before the segmentation head. + learning_rate: Adam learning rate. + use_face_grids: Fuse the face UV-grid geometry. + use_edge_grids: Fuse the edge UV-grid geometry. + ignore_index: Label value excluded from loss/metrics. + """ + + def __init__( + self, + num_classes: int, + face_feature_dim: int = 8, + edge_feature_dim: int = 7, + surf_emb_dim: int = 64, + curve_emb_dim: int = 64, + entity_hidden: int = 64, + hidden_dim: int = 128, + num_layers: int = 6, + dropout: float = 0.0, + learning_rate: float = 1e-3, + use_face_grids: bool = True, + use_edge_grids: bool = True, + ignore_index: int = IGNORE_INDEX, + ): + super().__init__() + self.save_hyperparameters() + + self.num_classes = num_classes + self.ignore_index = ignore_index + self.learning_rate = learning_rate + + self.surface_encoder = ( + UVNetSurfaceEncoder(in_channels=7, out_dim=surf_emb_dim) if use_face_grids else None + ) + self.curve_encoder = ( + UVNetCurveEncoder(in_channels=6, out_dim=curve_emb_dim) if use_edge_grids else None + ) + + face_in = face_feature_dim + (surf_emb_dim if use_face_grids else 0) + edge_in = edge_feature_dim + (curve_emb_dim if use_edge_grids else 0) + self.face_proj = nn.Sequential(nn.Linear(face_in, entity_hidden), nn.ReLU(inplace=True)) + self.edge_proj = nn.Sequential(nn.Linear(edge_in, entity_hidden), nn.ReLU(inplace=True)) + + coedge_in = 2 * entity_hidden + 1 # face_repr + edge_repr + reversed flag + self.encoder = BRepNetEncoder( + input_dim=coedge_in, + hidden_dim=hidden_dim, + output_dim=hidden_dim, + num_layers=num_layers, + ) + self.dropout = nn.Dropout(dropout) + self.seg_head = nn.Linear(hidden_dim, num_classes) + + metric_kwargs = {"num_classes": num_classes, "ignore_index": ignore_index} + self.val_iou = MulticlassJaccardIndex(average="macro", **metric_kwargs) + self.val_acc = MulticlassAccuracy(average="micro", **metric_kwargs) + self.test_iou = MulticlassJaccardIndex(average="macro", **metric_kwargs) + self.test_acc = MulticlassAccuracy(average="micro", **metric_kwargs) + + # -- forward ---------------------------------------------------------- + def forward(self, batch: BRepBatch) -> torch.Tensor: + """Return per-face class logits ``[num_faces, num_classes]``.""" + face_in = batch.face_features + if self.surface_encoder is not None: + face_geo = self.surface_encoder(batch.face_point_grids) + face_in = torch.cat([face_in, face_geo], dim=1) + face_repr = self.face_proj(face_in) + + edge_in = batch.edge_features + if self.curve_encoder is not None: + edge_geo = self.curve_encoder(batch.edge_point_grids) + edge_in = torch.cat([edge_in, edge_geo], dim=1) + edge_repr = self.edge_proj(edge_in) + + coedge_feats = torch.cat( + [ + face_repr[batch.coedge_to_face], + edge_repr[batch.coedge_to_edge], + batch.coedge_reversed, + ], + dim=1, + ) + coedge_data = CoedgeData( + features=coedge_feats, + next_indices=batch.coedge_to_next, + prev_indices=batch.coedge_to_prev, + mate_indices=batch.coedge_to_mate, + face_indices=batch.coedge_to_face, + ) + face_embeddings, _ = self.encoder(coedge_data) + return self.seg_head(self.dropout(face_embeddings)) + + # -- shared step ------------------------------------------------------ + def _compute_loss(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + if int((labels != self.ignore_index).sum()) == 0: + return logits.new_zeros((), requires_grad=True) + return F.cross_entropy(logits, labels, ignore_index=self.ignore_index) + + def training_step(self, batch: BRepBatch, batch_idx: int) -> torch.Tensor: + logits = self(batch) + loss = self._compute_loss(logits, batch.labels) + self.log("train_loss", loss, batch_size=batch.num_faces, prog_bar=True) + return loss + + def validation_step(self, batch: BRepBatch, batch_idx: int) -> torch.Tensor: + logits = self(batch) + loss = self._compute_loss(logits, batch.labels) + self.val_iou.update(logits, batch.labels) + self.val_acc.update(logits, batch.labels) + self.log("val_loss", loss, batch_size=batch.num_faces, prog_bar=True) + return loss + + def on_validation_epoch_end(self) -> None: + self.log("val_miou", self.val_iou.compute(), prog_bar=True) + self.log("val_acc", self.val_acc.compute(), prog_bar=True) + self.val_iou.reset() + self.val_acc.reset() + + def test_step(self, batch: BRepBatch, batch_idx: int) -> torch.Tensor: + logits = self(batch) + loss = self._compute_loss(logits, batch.labels) + self.test_iou.update(logits, batch.labels) + self.test_acc.update(logits, batch.labels) + self.log("test_loss", loss, batch_size=batch.num_faces) + return loss + + def on_test_epoch_end(self) -> None: + self.log("test_miou", self.test_iou.compute()) + self.log("test_acc", self.test_acc.compute()) + self.test_iou.reset() + self.test_acc.reset() + + def configure_optimizers(self): + return torch.optim.Adam(self.parameters(), lr=self.learning_rate) + + # -- inference helper ------------------------------------------------- + @torch.no_grad() + def predict_logits(self, batch: BRepBatch) -> list[tuple[str, torch.Tensor]]: + """Return ``(file_stem, per-face softmax probabilities)`` per solid.""" + self.eval() + probs = F.softmax(self(batch), dim=1) + out: list[tuple[str, torch.Tensor]] = [] + for stem, (start, end) in zip(batch.file_stems, batch.split_batch): + out.append((stem, probs[start:end].cpu())) + return out + + @staticmethod + def add_model_specific_args(parent_parser: argparse.ArgumentParser) -> argparse.ArgumentParser: + """Add ``LLBRepNet`` hyperparameters to an argparse parser.""" + parser = parent_parser.add_argument_group("LLBRepNet") + parser.add_argument("--surf-emb-dim", type=int, default=64) + parser.add_argument("--curve-emb-dim", type=int, default=64) + parser.add_argument("--entity-hidden", type=int, default=64) + parser.add_argument("--hidden-dim", type=int, default=128) + parser.add_argument("--num-layers", type=int, default=6) + parser.add_argument("--dropout", type=float, default=0.0) + parser.add_argument("--learning-rate", type=float, default=1e-3) + parser.add_argument("--no-face-grids", action="store_true") + parser.add_argument("--no-edge-grids", action="store_true") + return parent_parser diff --git a/ll_brepnet/ll_brepnet/models/uvnet_encoders.py b/ll_brepnet/ll_brepnet/models/uvnet_encoders.py index e69de29..11dc12c 100644 --- a/ll_brepnet/ll_brepnet/models/uvnet_encoders.py +++ b/ll_brepnet/ll_brepnet/models/uvnet_encoders.py @@ -0,0 +1,80 @@ +"""UV-grid geometry encoders (UV-Net style, written from scratch). + +Two small convolutional encoders turn the sampled B-Rep geometry into fixed-size +embeddings that are fused with the topological features: + +* :class:`UVNetSurfaceEncoder` -- a 2D CNN over a face's ``[7, U, V]`` UV-grid + (xyz + normal + trimming mask) -> ``[out_dim]``. +* :class:`UVNetCurveEncoder` -- a 1D CNN over an edge's ``[6, U]`` U-grid + (xyz + tangent) -> ``[out_dim]``. + +These are standard convolutional stacks (the generic UV-Net idea of running a +CNN over the parametric grid), implemented independently for ``ll_brepnet``. +""" + +from __future__ import annotations + +import torch +from torch import nn + + +class UVNetSurfaceEncoder(nn.Module): + """2D CNN encoder for face UV-grids ``[B, in_channels, U, V]``. + + Args: + in_channels: Input channels (7 = xyz + normal + trimming mask). + out_dim: Output embedding dimension. + """ + + def __init__(self, in_channels: int = 7, out_dim: int = 64) -> None: + super().__init__() + self.net = nn.Sequential( + nn.Conv2d(in_channels, 32, kernel_size=3, padding=1), + nn.BatchNorm2d(32), + nn.ReLU(inplace=True), + nn.Conv2d(32, 64, kernel_size=3, padding=1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True), + nn.Conv2d(64, out_dim, kernel_size=3, padding=1), + nn.BatchNorm2d(out_dim), + nn.ReLU(inplace=True), + nn.AdaptiveAvgPool2d(1), + ) + self.out_dim = out_dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """``[B, in_channels, U, V] -> [B, out_dim]``.""" + if x.shape[0] == 0: + return x.new_zeros((0, self.out_dim)) + return self.net(x).flatten(1) + + +class UVNetCurveEncoder(nn.Module): + """1D CNN encoder for edge U-grids ``[B, in_channels, U]``. + + Args: + in_channels: Input channels (6 = xyz + tangent). + out_dim: Output embedding dimension. + """ + + def __init__(self, in_channels: int = 6, out_dim: int = 64) -> None: + super().__init__() + self.net = nn.Sequential( + nn.Conv1d(in_channels, 32, kernel_size=3, padding=1), + nn.BatchNorm1d(32), + nn.ReLU(inplace=True), + nn.Conv1d(32, 64, kernel_size=3, padding=1), + nn.BatchNorm1d(64), + nn.ReLU(inplace=True), + nn.Conv1d(64, out_dim, kernel_size=3, padding=1), + nn.BatchNorm1d(out_dim), + nn.ReLU(inplace=True), + nn.AdaptiveAvgPool1d(1), + ) + self.out_dim = out_dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """``[B, in_channels, U] -> [B, out_dim]``.""" + if x.shape[0] == 0: + return x.new_zeros((0, self.out_dim)) + return self.net(x).flatten(1) diff --git a/ll_brepnet/ll_brepnet/pipelines/__init__.py b/ll_brepnet/ll_brepnet/pipelines/__init__.py new file mode 100644 index 0000000..02d0c35 --- /dev/null +++ b/ll_brepnet/ll_brepnet/pipelines/__init__.py @@ -0,0 +1,4 @@ +"""Data pipelines: STEP/JSON -> coedge graph + geometry -> ``.npz`` extraction +and train/val/test dataset-manifest building.""" + +from __future__ import annotations diff --git a/ll_brepnet/ll_brepnet/pipelines/build_dataset_file.py b/ll_brepnet/ll_brepnet/pipelines/build_dataset_file.py index e69de29..579c61a 100644 --- a/ll_brepnet/ll_brepnet/pipelines/build_dataset_file.py +++ b/ll_brepnet/ll_brepnet/pipelines/build_dataset_file.py @@ -0,0 +1,215 @@ +"""Build the dataset manifest (``dataset.json``) for a folder of ``.npz`` records. + +The manifest defines the train / validation / test split and the per-feature +standardisation statistics (mean + standard deviation) computed over the +**training split only**. The dataset loader applies these to z-score the +continuous features at load time, so statistics never leak from val/test. + +Manifest schema:: + + { + "training_set": ["", ...], + "validation_set": ["", ...], + "test_set": ["", ...], + "feature_standardization": { + "face_features": [{"mean": m, "standard_deviation": s}, ...], + "edge_features": [{"mean": m, "standard_deviation": s}, ...] + }, + "num_classes": , + "class_names": [, ...] # optional + } +""" + +from __future__ import annotations + +import argparse +import json +import logging +from collections.abc import Sequence +from pathlib import Path + +import numpy as np +from sklearn.model_selection import train_test_split + +_log = logging.getLogger(__name__) + +# Features whose per-column mean/std the loader will use to z-score. +STANDARDIZED_FEATURES = ("face_features", "edge_features") + + +def _stems_in(npz_dir: Path) -> list[str]: + return sorted(p.stem for p in Path(npz_dir).glob("*.npz")) + + +def compute_standardization( + npz_dir: Path, + train_stems: Sequence[str], + feature_keys: Sequence[str] = STANDARDIZED_FEATURES, +) -> dict[str, list[dict[str, float]]]: + """Compute per-column mean/std for ``feature_keys`` over the training set. + + Uses a single streaming pass (sum + sum-of-squares) so it scales to large + datasets. A near-zero standard deviation is floored to 1.0 so the loader's + division is well-defined for constant columns (e.g. an always-off one-hot). + """ + npz_dir = Path(npz_dir) + accum: dict[str, dict[str, np.ndarray]] = {} + + for stem in train_stems: + path = npz_dir / f"{stem}.npz" + if not path.exists(): + _log.warning("Skipping missing npz for standardization: %s", path.name) + continue + with np.load(path) as data: + for key in feature_keys: + if key not in data: + continue + arr = np.asarray(data[key], dtype=np.float64) + if arr.ndim != 2 or arr.shape[0] == 0: + continue + if key not in accum: + ncol = arr.shape[1] + accum[key] = { + "sum": np.zeros(ncol, dtype=np.float64), + "sumsq": np.zeros(ncol, dtype=np.float64), + "count": np.zeros(1, dtype=np.float64), + } + accum[key]["sum"] += arr.sum(axis=0) + accum[key]["sumsq"] += (arr**2).sum(axis=0) + accum[key]["count"][0] += arr.shape[0] + + standardization: dict[str, list[dict[str, float]]] = {} + for key, acc in accum.items(): + count = max(float(acc["count"][0]), 1.0) + mean = acc["sum"] / count + var = np.maximum(acc["sumsq"] / count - mean**2, 0.0) + std = np.sqrt(var) + std[std < 1e-6] = 1.0 + standardization[key] = [ + {"mean": float(m), "standard_deviation": float(s)} for m, s in zip(mean, std) + ] + return standardization + + +def build_dataset_file( + npz_dir: Path, + output_file: Path, + validation_split: float = 0.2, + test_split: float = 0.15, + train_test_file: Path | None = None, + class_names: Sequence[str] | None = None, + seed: int = 42, +) -> dict: + """Create train/val/test splits + standardization and write ``dataset.json``. + + Args: + npz_dir: Directory of ``.npz`` records. + output_file: Path to write the manifest JSON to. + validation_split: Fraction of the (post-test) training pool held out for + validation. + test_split: Fraction of all solids held out for test (ignored when + ``train_test_file`` is given). + train_test_file: Optional JSON ``{"train": [...], "test": [...]}`` of + stems to use instead of a random test split. + class_names: Optional ordered list of segment class names. + seed: RNG seed for the random splits (deterministic). + + Returns: + The manifest dictionary that was written. + """ + npz_dir = Path(npz_dir) + all_stems = _stems_in(npz_dir) + if not all_stems: + raise ValueError(f"No .npz records found in {npz_dir}") + + if train_test_file is not None: + split = json.loads(Path(train_test_file).read_text()) + test_stems = [s for s in split.get("test", []) if s in set(all_stems)] + train_pool = [s for s in split.get("train", []) if s in set(all_stems)] + if not train_pool: + train_pool = [s for s in all_stems if s not in set(test_stems)] + elif test_split and test_split > 0.0 and len(all_stems) > 2: + train_pool, test_stems = train_test_split( + all_stems, test_size=test_split, random_state=seed + ) + else: + train_pool, test_stems = all_stems, [] + + if validation_split and validation_split > 0.0 and len(train_pool) > 2: + train_stems, val_stems = train_test_split( + train_pool, test_size=validation_split, random_state=seed + ) + else: + train_stems, val_stems = train_pool, [] + + train_stems, val_stems, test_stems = ( + sorted(train_stems), + sorted(val_stems), + sorted(test_stems), + ) + + standardization = compute_standardization(npz_dir, train_stems) + + num_classes: int | None = len(class_names) if class_names else None + + manifest: dict = { + "training_set": train_stems, + "validation_set": val_stems, + "test_set": test_stems, + "feature_standardization": standardization, + "num_classes": num_classes, + "class_names": list(class_names) if class_names else [], + } + + output_file = Path(output_file) + output_file.parent.mkdir(parents=True, exist_ok=True) + output_file.write_text(json.dumps(manifest, indent=2)) + _log.info( + "Wrote %s: train=%d val=%d test=%d", + output_file, + len(train_stems), + len(val_stems), + len(test_stems), + ) + return manifest + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Build a BRepNet dataset manifest.") + parser.add_argument("--npz-dir", required=True, help="Directory of .npz records") + parser.add_argument("--output", required=True, help="Path to write dataset.json") + parser.add_argument("--validation-split", type=float, default=0.2) + parser.add_argument("--test-split", type=float, default=0.15) + parser.add_argument("--train-test-file", default=None) + parser.add_argument( + "--class-names-file", + default=None, + help="Optional JSON list of segment class names", + ) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args(argv) + + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") + class_names = None + if args.class_names_file: + class_names = json.loads(Path(args.class_names_file).read_text()) + + manifest = build_dataset_file( + Path(args.npz_dir), + Path(args.output), + validation_split=args.validation_split, + test_split=args.test_split, + train_test_file=Path(args.train_test_file) if args.train_test_file else None, + class_names=class_names, + seed=args.seed, + ) + print( + f"train={len(manifest['training_set'])} " + f"val={len(manifest['validation_set'])} " + f"test={len(manifest['test_set'])}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ll_brepnet/ll_brepnet/pipelines/entity_mapper.py b/ll_brepnet/ll_brepnet/pipelines/entity_mapper.py index e69de29..ea24d90 100644 --- a/ll_brepnet/ll_brepnet/pipelines/entity_mapper.py +++ b/ll_brepnet/ll_brepnet/pipelines/entity_mapper.py @@ -0,0 +1,114 @@ +"""Stable integer indices for the faces and edges of a B-Rep solid. + +``BRepEntityMapper`` assigns each face and edge of a ``TopoDS_Shape`` a stable +0-based index and keeps the inverse (index -> ``TopoDS`` entity) so that +per-face / per-edge feature and UV-grid arrays can be built in a consistent +order. Coedges produced by ``cadling``'s ``CoedgeExtractor`` carry string +``face_id`` / ``edge_id`` identifiers; this mapper computes the **same** stable +identifier (OCC ``HashCode`` with a ``hash()`` fallback) so a coedge can be +resolved to the index of its parent face/edge. + +This module is deliberately self-contained: it does not rely on +``cadling``'s ``ShapeIdentityRegistry``, whose ``register_all_*`` helpers can +silently register nothing on some pythonocc builds (a bad ``topods`` import), +which would leave every ``face_index`` at ``-1``. +""" + +from __future__ import annotations + +import logging + +_log = logging.getLogger(__name__) + +HAS_OCC = False +try: + from OCC.Core.TopAbs import TopAbs_EDGE, TopAbs_FACE + from OCC.Core.TopExp import TopExp_Explorer + from OCC.Core.TopoDS import TopoDS_Edge, TopoDS_Face, TopoDS_Shape, topods + + HAS_OCC = True +except ImportError: # pragma: no cover - exercised only without pythonocc + _log.debug("pythonocc-core not available; BRepEntityMapper is inert.") + + +def stable_shape_id(shape: TopoDS_Shape) -> str: + """Return a stable identifier for an OCC shape. + + Uses OCCT ``HashCode`` when exposed (stable across the different Python + wrapper objects OCC hands back for the same topological entity), falling + back to Python's ``hash`` on pythonocc builds that drop ``HashCode``. This + matches ``cadling.lib.topology.face_identity.ShapeIdentityRegistry.get_id`` + so coedge ``face_id`` / ``edge_id`` strings line up with this mapper. + """ + try: + return str(shape.HashCode(2**31 - 1)) + except (AttributeError, TypeError): + return str(hash(shape)) + + +class BRepEntityMapper: + """Assigns stable 0-based indices to the faces and edges of a solid. + + Args: + shape: The ``TopoDS_Shape`` (typically a solid) to index. + """ + + def __init__(self, shape: TopoDS_Shape): + self._faces: list[TopoDS_Face] = [] + self._edges: list[TopoDS_Edge] = [] + self._face_id_to_index: dict[str, int] = {} + self._edge_id_to_index: dict[str, int] = {} + if HAS_OCC and shape is not None: + self._build(shape) + + def _build(self, shape: TopoDS_Shape) -> None: + face_exp = TopExp_Explorer(shape, TopAbs_FACE) + while face_exp.More(): + face = topods.Face(face_exp.Current()) + fid = stable_shape_id(face) + if fid not in self._face_id_to_index: + self._face_id_to_index[fid] = len(self._faces) + self._faces.append(face) + face_exp.Next() + + edge_exp = TopExp_Explorer(shape, TopAbs_EDGE) + while edge_exp.More(): + edge = topods.Edge(edge_exp.Current()) + eid = stable_shape_id(edge) + if eid not in self._edge_id_to_index: + self._edge_id_to_index[eid] = len(self._edges) + self._edges.append(edge) + edge_exp.Next() + + # -- counts ----------------------------------------------------------- + @property + def num_faces(self) -> int: + return len(self._faces) + + @property + def num_edges(self) -> int: + return len(self._edges) + + # -- index -> entity -------------------------------------------------- + def face_by_index(self, index: int) -> TopoDS_Face: + return self._faces[index] + + def edge_by_index(self, index: int) -> TopoDS_Edge: + return self._edges[index] + + def faces(self) -> list[TopoDS_Face]: + """Faces in index order.""" + return list(self._faces) + + def edges(self) -> list[TopoDS_Edge]: + """Edges in index order.""" + return list(self._edges) + + # -- id -> index ------------------------------------------------------ + def face_index(self, face_id: str) -> int | None: + """Index of the face with the given stable id, or ``None``.""" + return self._face_id_to_index.get(face_id) + + def edge_index(self, edge_id: str) -> int | None: + """Index of the edge with the given stable id, or ``None``.""" + return self._edge_id_to_index.get(edge_id) diff --git a/ll_brepnet/ll_brepnet/pipelines/extract_brepnet_data_from_json.py b/ll_brepnet/ll_brepnet/pipelines/extract_brepnet_data_from_json.py index e69de29..9476b90 100644 --- a/ll_brepnet/ll_brepnet/pipelines/extract_brepnet_data_from_json.py +++ b/ll_brepnet/ll_brepnet/pipelines/extract_brepnet_data_from_json.py @@ -0,0 +1,168 @@ +"""Build a BRepNet-style ``.npz`` record from a precomputed JSON description. + +This is the non-STEP front-end: when a caller already has the B-Rep coedge +topology and per-entity features (e.g. exported from another CAD tool or a +cached intermediate), they can supply it as JSON and obtain a ``.npz`` record +identical in schema to :mod:`ll_brepnet.pipelines.extract_brepnet_data_from_step`, +ready for the dataset loader. + +Expected JSON schema (UV-grids are optional and default to zeros):: + + { + "coedge_to_next": [int, ...], # length C + "coedge_to_prev": [int, ...], # length C + "coedge_to_mate": [int, ...], # length C + "coedge_to_face": [int, ...], # length C, values in [0, F) + "coedge_to_edge": [int, ...], # length C, values in [0, E) + "coedge_reversed": [0|1, ...], # optional, length C, default zeros + "face_features": [[float, ...], ...], # [F, Df] + "edge_features": [[float, ...], ...], # [E, De] + "face_point_grids": [...], # optional [F, 7, U, V] + "edge_point_grids": [...] # optional [E, 6, U] + } +""" + +from __future__ import annotations + +import argparse +import json +import logging +from pathlib import Path + +import numpy as np + +from .extract_brepnet_data_from_step import NUM_U, NUM_V + +_log = logging.getLogger(__name__) + +_COEDGE_INT_KEYS = ( + "coedge_to_next", + "coedge_to_prev", + "coedge_to_mate", + "coedge_to_face", + "coedge_to_edge", +) + + +class BRepJsonExtractor: + """Convert a JSON topology/feature description into BRepNet ``.npz`` arrays. + + Args: + topology: The parsed JSON dictionary (see module docstring for schema). + num_u, num_v: UV-grid resolution used when grids are absent. + """ + + def __init__(self, topology: dict, num_u: int = NUM_U, num_v: int = NUM_V): + self.topology = topology + self.num_u = num_u + self.num_v = num_v + + def extract_arrays(self) -> dict[str, np.ndarray]: + """Validate the description and return the named ``.npz`` arrays.""" + t = self.topology + + for key in (*_COEDGE_INT_KEYS, "face_features", "edge_features"): + if key not in t: + raise ValueError(f"JSON topology missing required key: {key!r}") + + face_features = np.asarray(t["face_features"], dtype=np.float32) + edge_features = np.asarray(t["edge_features"], dtype=np.float32) + if face_features.ndim != 2 or edge_features.ndim != 2: + raise ValueError("face_features and edge_features must be 2-D [N, D]") + num_faces, num_edges = face_features.shape[0], edge_features.shape[0] + + coedge_arrays: dict[str, np.ndarray] = {} + lengths = set() + for key in _COEDGE_INT_KEYS: + arr = np.asarray(t[key], dtype=np.int64).ravel() + coedge_arrays[key] = arr + lengths.add(arr.shape[0]) + if len(lengths) != 1: + raise ValueError(f"coedge_to_* arrays have inconsistent lengths: {lengths}") + num_coedges = lengths.pop() + + # Bounds checks: incidence must reference valid entities. + if num_coedges: + if int(coedge_arrays["coedge_to_face"].max(initial=0)) >= num_faces: + raise ValueError("coedge_to_face references a face index >= num_faces") + if int(coedge_arrays["coedge_to_edge"].max(initial=0)) >= num_edges: + raise ValueError("coedge_to_edge references an edge index >= num_edges") + for key in ("coedge_to_next", "coedge_to_prev", "coedge_to_mate"): + if int(coedge_arrays[key].max(initial=0)) >= num_coedges: + raise ValueError(f"{key} references a coedge index >= num_coedges") + + if "coedge_reversed" in t: + coedge_reversed = np.asarray(t["coedge_reversed"], dtype=np.float32).ravel() + if coedge_reversed.shape[0] != num_coedges: + raise ValueError("coedge_reversed length must equal num_coedges") + else: + coedge_reversed = np.zeros(num_coedges, dtype=np.float32) + + if "face_point_grids" in t: + face_point_grids = np.asarray(t["face_point_grids"], dtype=np.float32) + if face_point_grids.shape != (num_faces, 7, self.num_u, self.num_v): + raise ValueError( + "face_point_grids must have shape " + f"[{num_faces}, 7, {self.num_u}, {self.num_v}]" + ) + else: + face_point_grids = np.zeros((num_faces, 7, self.num_u, self.num_v), dtype=np.float32) + + if "edge_point_grids" in t: + edge_point_grids = np.asarray(t["edge_point_grids"], dtype=np.float32) + if edge_point_grids.shape != (num_edges, 6, self.num_u): + raise ValueError(f"edge_point_grids must have shape [{num_edges}, 6, {self.num_u}]") + else: + edge_point_grids = np.zeros((num_edges, 6, self.num_u), dtype=np.float32) + + return { + **coedge_arrays, + "coedge_reversed": coedge_reversed, + "face_features": face_features, + "edge_features": edge_features, + "face_point_grids": face_point_grids, + "edge_point_grids": edge_point_grids, + "num_faces": np.int64(num_faces), + "num_edges": np.int64(num_edges), + "num_coedges": np.int64(num_coedges), + } + + def process(self, output_path: Path) -> Path: + """Extract arrays and write them to ``output_path`` (a ``.npz`` file).""" + arrays = self.extract_arrays() + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed(output_path, **arrays) + _log.info( + "Wrote %s (F=%d E=%d C=%d)", + output_path.name, + int(arrays["num_faces"]), + int(arrays["num_edges"]), + int(arrays["num_coedges"]), + ) + return output_path + + +def extract_brepnet_data_from_json(json_path: Path, output_path: Path) -> Path: + """Read a JSON topology file and write the corresponding ``.npz`` record.""" + json_path = Path(json_path) + topology = json.loads(json_path.read_text()) + return BRepJsonExtractor(topology).process(output_path) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description="Build a BRepNet-style .npz from a JSON topology description." + ) + parser.add_argument("--json", required=True, help="Input JSON topology file") + parser.add_argument("--output", required=True, help="Output .npz path") + args = parser.parse_args(argv) + + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") + out = extract_brepnet_data_from_json(Path(args.json), Path(args.output)) + print(f"Wrote {out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ll_brepnet/ll_brepnet/pipelines/extract_brepnet_data_from_step.py b/ll_brepnet/ll_brepnet/pipelines/extract_brepnet_data_from_step.py index e69de29..c4eb274 100644 --- a/ll_brepnet/ll_brepnet/pipelines/extract_brepnet_data_from_step.py +++ b/ll_brepnet/ll_brepnet/pipelines/extract_brepnet_data_from_step.py @@ -0,0 +1,499 @@ +"""Extract a BRepNet-style ``.npz`` training record from a STEP file. + +The record captures, for one solid: + +* the **coedge graph** -- ``coedge_to_next / prev / mate / face / edge`` integer + incidence arrays (oriented half-edge adjacency), plus a ``coedge_reversed`` + flag; +* per-**face** features (surface-type one-hot + area) and a ``[7, U, V]`` UV-grid + (xyz + normal + trimming mask); +* per-**edge** features (curve-type one-hot + length + convexity) and a + ``[6, U]`` U-grid (xyz + tangent). + +Geometry is normalised to the unit box ``[-1, 1]^3`` before sampling. The coedge +topology is obtained from ``cadling``'s ``CoedgeExtractor`` (oriented loops via +``BRepTools_WireExplorer`` + mate finding via ``MapShapesAndAncestors``); all +face/edge indexing, features and UV-grids are computed here so the record does +not depend on ``cadling``'s shape-identity registry (which can register nothing +on some pythonocc builds). + +This is an independent, MIT-licensed implementation -- see ``ATTRIBUTION.md``. +""" + +from __future__ import annotations + +import argparse +import logging +from pathlib import Path + +import numpy as np + +from .entity_mapper import BRepEntityMapper + +_log = logging.getLogger(__name__) + +# UV-grid resolution (UV-Net default). +NUM_U = 10 +NUM_V = 10 + +# Feature layouts (kept small and explicit; standardisation happens later in the +# dataset using statistics over the training split). +SURFACE_TYPES = ["plane", "cylinder", "cone", "sphere", "torus", "bspline", "other"] +CURVE_TYPES = ["line", "circle", "ellipse", "bspline", "other"] +NUM_FACE_FEATURES = len(SURFACE_TYPES) + 1 # + area +NUM_EDGE_FEATURES = len(CURVE_TYPES) + 2 # + length + convexity + +HAS_OCC = False +try: + from cadling.lib.topology.coedge_extractor import CoedgeExtractor + from OCC.Core.Bnd import Bnd_Box + from OCC.Core.BRep import BRep_Tool + from OCC.Core.BRepAdaptor import BRepAdaptor_Curve, BRepAdaptor_Surface + from OCC.Core.BRepBndLib import brepbndlib + from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform + from OCC.Core.BRepGProp import brepgprop + from OCC.Core.BRepLProp import BRepLProp_SLProps + from OCC.Core.GeomAbs import ( + GeomAbs_BezierCurve, + GeomAbs_BezierSurface, + GeomAbs_BSplineCurve, + GeomAbs_BSplineSurface, + GeomAbs_Circle, + GeomAbs_Cone, + GeomAbs_Cylinder, + GeomAbs_Ellipse, + GeomAbs_Line, + GeomAbs_Plane, + GeomAbs_Sphere, + GeomAbs_Torus, + ) + from OCC.Core.GeomAPI import GeomAPI_ProjectPointOnSurf + from OCC.Core.gp import gp_Pnt, gp_Trsf, gp_Vec + from OCC.Core.GProp import GProp_GProps + from OCC.Core.IFSelect import IFSelect_RetDone + from OCC.Core.STEPControl import STEPControl_Reader + from OCC.Core.TopoDS import TopoDS_Edge, TopoDS_Face, TopoDS_Shape + from occwl.edge import Edge as OCCWLEdge + from occwl.face import Face as OCCWLFace + from occwl.uvgrid import ugrid, uvgrid + + HAS_OCC = True +except ImportError as _exc: # pragma: no cover - exercised only without pythonocc + _log.debug("pythonocc/occwl/cadling not available; STEP extraction is inert: %s", _exc) + + +# --------------------------------------------------------------------------- +# STEP loading + normalisation +# --------------------------------------------------------------------------- + + +def load_step_shape(step_file: Path) -> TopoDS_Shape: + """Read a STEP file and return its (compound) ``TopoDS_Shape``.""" + reader = STEPControl_Reader() + status = reader.ReadFile(str(step_file)) + if status != IFSelect_RetDone: + raise OSError(f"Failed to read STEP file: {step_file}") + reader.TransferRoots() + return reader.OneShape() + + +def scale_shape_to_unit_box(shape: TopoDS_Shape) -> TopoDS_Shape: + """Center ``shape`` at the origin and uniformly scale it into ``[-1, 1]^3``. + + A single uniform scale (the largest half-extent maps to 1.0) preserves the + aspect ratio, matching how UV-Net / BRepNet normalise solids before + sampling geometry. + """ + box = Bnd_Box() + brepbndlib.Add(shape, box) + xmin, ymin, zmin, xmax, ymax, zmax = box.Get() + cx, cy, cz = (xmin + xmax) / 2.0, (ymin + ymax) / 2.0, (zmin + zmax) / 2.0 + half = max(xmax - xmin, ymax - ymin, zmax - zmin) / 2.0 + if half <= 0.0: + return shape + scale = 1.0 / half + + trsf = gp_Trsf() + trsf.SetTranslation(gp_Vec(-cx, -cy, -cz)) + scale_trsf = gp_Trsf() + scale_trsf.SetScale(gp_Pnt(0.0, 0.0, 0.0), scale) + combined = scale_trsf.Multiplied(trsf) + return BRepBuilderAPI_Transform(shape, combined, True).Shape() + + +# --------------------------------------------------------------------------- +# Per-face / per-edge geometry features +# --------------------------------------------------------------------------- + + +def _one_hot(index: int, size: int) -> np.ndarray: + vec = np.zeros(size, dtype=np.float32) + vec[index] = 1.0 + return vec + + +def face_surface_onehot(face: TopoDS_Face) -> np.ndarray: + """One-hot of the face's surface type over :data:`SURFACE_TYPES`.""" + surf = BRepAdaptor_Surface(face) + t = surf.GetType() + mapping = { + GeomAbs_Plane: 0, + GeomAbs_Cylinder: 1, + GeomAbs_Cone: 2, + GeomAbs_Sphere: 3, + GeomAbs_Torus: 4, + GeomAbs_BSplineSurface: 5, + GeomAbs_BezierSurface: 5, + } + return _one_hot(mapping.get(t, len(SURFACE_TYPES) - 1), len(SURFACE_TYPES)) + + +def face_area(face: TopoDS_Face) -> float: + props = GProp_GProps() + brepgprop.SurfaceProperties(face, props) + return float(props.Mass()) + + +def edge_curve_onehot(edge: TopoDS_Edge) -> np.ndarray: + """One-hot of the edge's curve type over :data:`CURVE_TYPES`.""" + try: + curve = BRepAdaptor_Curve(edge) + t = curve.GetType() + except Exception: + return _one_hot(len(CURVE_TYPES) - 1, len(CURVE_TYPES)) + mapping = { + GeomAbs_Line: 0, + GeomAbs_Circle: 1, + GeomAbs_Ellipse: 2, + GeomAbs_BSplineCurve: 3, + GeomAbs_BezierCurve: 3, + } + return _one_hot(mapping.get(t, len(CURVE_TYPES) - 1), len(CURVE_TYPES)) + + +def edge_length(edge: TopoDS_Edge) -> float: + props = GProp_GProps() + brepgprop.LinearProperties(edge, props) + return float(props.Mass()) + + +def _face_normal_at(face: TopoDS_Face, pnt: gp_Pnt) -> np.ndarray | None: + """Outward unit normal of ``face`` at the surface point nearest ``pnt``.""" + surface = BRep_Tool.Surface(face) + proj = GeomAPI_ProjectPointOnSurf(pnt, surface) + if proj.NbPoints() < 1: + return None + u, v = proj.LowerDistanceParameters() + props = BRepLProp_SLProps(BRepAdaptor_Surface(face), float(u), float(v), 1, 1e-6) + if not props.IsNormalDefined(): + return None + n = props.Normal() + normal = np.array([n.X(), n.Y(), n.Z()], dtype=np.float64) + norm = np.linalg.norm(normal) + if norm < 1e-12: + return None + normal /= norm + # Account for the face orientation flag (the surface normal is geometric). + from OCC.Core.TopAbs import TopAbs_REVERSED + + if face.Orientation() == TopAbs_REVERSED: + normal = -normal + return normal + + +def _edge_mid_point_tangent(edge: TopoDS_Edge) -> tuple[gp_Pnt, np.ndarray] | None: + """Return the 3D midpoint and unit tangent of ``edge``.""" + try: + curve = BRepAdaptor_Curve(edge) + u_mid = 0.5 * (curve.FirstParameter() + curve.LastParameter()) + pnt = gp_Pnt() + vec = gp_Vec() + curve.D1(u_mid, pnt, vec) + tangent = np.array([vec.X(), vec.Y(), vec.Z()], dtype=np.float64) + norm = np.linalg.norm(tangent) + if norm < 1e-12: + return None + return pnt, tangent / norm + except Exception: + return None + + +def edge_convexity( + edge: TopoDS_Edge, + face_a: TopoDS_Face, + face_b: TopoDS_Face, +) -> float: + """Signed convexity of an edge shared by ``face_a`` and ``face_b``. + + Computes ``s = (n_a x t) . n_b`` at the edge midpoint, where ``n_a`` / + ``n_b`` are the (orientation-corrected) face normals and ``t`` the edge + tangent. ``s < 0`` -> convex (1.0), ``s > 0`` -> concave (0.0), and + near-tangent edges (``|s|`` small) -> smooth (0.5). Returns 0.5 when the + geometry cannot be evaluated (e.g. a boundary/seam edge). + """ + mt = _edge_mid_point_tangent(edge) + if mt is None: + return 0.5 + pnt, tangent = mt + n_a = _face_normal_at(face_a, pnt) + n_b = _face_normal_at(face_b, pnt) + if n_a is None or n_b is None: + return 0.5 + s = float(np.dot(np.cross(n_a, tangent), n_b)) + if abs(s) < 1e-4: + return 0.5 + return 1.0 if s < 0.0 else 0.0 + + +# --------------------------------------------------------------------------- +# UV-grids (point + normal + inside mask for faces; point + tangent for edges) +# --------------------------------------------------------------------------- + + +def face_uv_grid(face: TopoDS_Face, num_u: int = NUM_U, num_v: int = NUM_V) -> np.ndarray: + """Return a ``[7, num_u, num_v]`` UV-grid: xyz(3) + normal(3) + mask(1). + + Channels-first so it feeds directly into a 2D CNN surface encoder. Failed + samples yield a zero grid (the trimming mask stays 0), which the model can + distinguish from material via that mask channel. + """ + try: + wrapped = OCCWLFace(face) + points = uvgrid(wrapped, num_u=num_u, num_v=num_v, method="point") + normals = uvgrid(wrapped, num_u=num_u, num_v=num_v, method="normal") + inside = uvgrid(wrapped, num_u=num_u, num_v=num_v, method="inside") + if points is None or normals is None or inside is None: + return np.zeros((7, num_u, num_v), dtype=np.float32) + grid = np.concatenate([points, normals, inside.astype(np.float32)], axis=2) # [u, v, 7] + return np.transpose(grid, (2, 0, 1)).astype(np.float32) + except Exception as exc: + _log.debug("face UV-grid extraction failed: %s", exc) + return np.zeros((7, num_u, num_v), dtype=np.float32) + + +def edge_u_grid(edge: TopoDS_Edge, num_u: int = NUM_U) -> np.ndarray: + """Return a ``[6, num_u]`` U-grid: xyz(3) + tangent(3), channels-first.""" + try: + wrapped = OCCWLEdge(edge) + points = ugrid(wrapped, num_u=num_u, method="point") + tangents = ugrid(wrapped, num_u=num_u, method="tangent") + if points is None or tangents is None: + return np.zeros((6, num_u), dtype=np.float32) + grid = np.concatenate([points, tangents], axis=1) # [u, 6] + return np.transpose(grid, (1, 0)).astype(np.float32) + except Exception as exc: + _log.debug("edge U-grid extraction failed: %s", exc) + return np.zeros((6, num_u), dtype=np.float32) + + +# --------------------------------------------------------------------------- +# Extractor +# --------------------------------------------------------------------------- + + +class BRepDataExtractor: + """Extract BRepNet-style arrays from one STEP file. + + Args: + step_file: Path to the ``.step`` / ``.stp`` file. + scale_body: Normalise the solid into the unit box before sampling. + num_u, num_v: UV-grid resolution. + """ + + def __init__( + self, + step_file: Path, + scale_body: bool = True, + num_u: int = NUM_U, + num_v: int = NUM_V, + ): + self.step_file = Path(step_file) + self.scale_body = scale_body + self.num_u = num_u + self.num_v = num_v + + def extract_arrays(self) -> dict[str, np.ndarray]: + """Return the full set of named arrays for this solid.""" + if not HAS_OCC: + raise RuntimeError("pythonocc-core / occwl / cadling are required for STEP extraction") + + shape = load_step_shape(self.step_file) + if self.scale_body: + shape = scale_shape_to_unit_box(shape) + + mapper = BRepEntityMapper(shape) + num_faces = mapper.num_faces + num_edges = mapper.num_edges + if num_faces == 0 or num_edges == 0: + raise ValueError(f"No faces/edges extracted from {self.step_file}") + + coedges = CoedgeExtractor().extract_coedges(shape) + if not coedges: + raise ValueError(f"No coedges extracted from {self.step_file}") + num_coedges = len(coedges) + + # Coedge id -> position so next/prev/mate (which reference coedge ids) + # become 0-based array indices. + id_to_pos = {c.id: pos for pos, c in enumerate(coedges)} + + coedge_to_next = np.arange(num_coedges, dtype=np.int64) + coedge_to_prev = np.arange(num_coedges, dtype=np.int64) + coedge_to_mate = np.arange(num_coedges, dtype=np.int64) + coedge_to_face = np.zeros(num_coedges, dtype=np.int64) + coedge_to_edge = np.zeros(num_coedges, dtype=np.int64) + coedge_reversed = np.zeros(num_coedges, dtype=np.float32) + + for pos, c in enumerate(coedges): + if c.next_id is not None: + coedge_to_next[pos] = id_to_pos.get(c.next_id, pos) + if c.prev_id is not None: + coedge_to_prev[pos] = id_to_pos.get(c.prev_id, pos) + if c.mate_id is not None: + coedge_to_mate[pos] = id_to_pos.get(c.mate_id, pos) + fi = mapper.face_index(c.face_id) + ei = mapper.edge_index(c.edge_id) + coedge_to_face[pos] = fi if fi is not None else 0 + coedge_to_edge[pos] = ei if ei is not None else 0 + coedge_reversed[pos] = 1.0 if c.orientation == "REVERSED" else 0.0 + + # Per-face features + grids (index order from the mapper). + face_features = np.zeros((num_faces, NUM_FACE_FEATURES), dtype=np.float32) + face_point_grids = np.zeros((num_faces, 7, self.num_u, self.num_v), dtype=np.float32) + for i in range(num_faces): + face = mapper.face_by_index(i) + face_features[i, : len(SURFACE_TYPES)] = face_surface_onehot(face) + face_features[i, len(SURFACE_TYPES)] = face_area(face) + face_point_grids[i] = face_uv_grid(face, self.num_u, self.num_v) + + # Adjacent faces per edge (for convexity), from the coedge incidence. + edge_faces: dict[int, list[int]] = {} + for pos in range(num_coedges): + edge_faces.setdefault(int(coedge_to_edge[pos]), []) + fi = int(coedge_to_face[pos]) + if fi not in edge_faces[int(coedge_to_edge[pos])]: + edge_faces[int(coedge_to_edge[pos])].append(fi) + + # Per-edge features + grids. + edge_features = np.zeros((num_edges, NUM_EDGE_FEATURES), dtype=np.float32) + edge_point_grids = np.zeros((num_edges, 6, self.num_u), dtype=np.float32) + for j in range(num_edges): + edge = mapper.edge_by_index(j) + edge_features[j, : len(CURVE_TYPES)] = edge_curve_onehot(edge) + edge_features[j, len(CURVE_TYPES)] = edge_length(edge) + faces_j = edge_faces.get(j, []) + if len(faces_j) >= 2: + conv = edge_convexity( + edge, mapper.face_by_index(faces_j[0]), mapper.face_by_index(faces_j[1]) + ) + else: + conv = 0.5 + edge_features[j, len(CURVE_TYPES) + 1] = conv + edge_point_grids[j] = edge_u_grid(edge, self.num_u) + + return { + "coedge_to_next": coedge_to_next, + "coedge_to_prev": coedge_to_prev, + "coedge_to_mate": coedge_to_mate, + "coedge_to_face": coedge_to_face, + "coedge_to_edge": coedge_to_edge, + "coedge_reversed": coedge_reversed, + "face_features": face_features, + "face_point_grids": face_point_grids, + "edge_features": edge_features, + "edge_point_grids": edge_point_grids, + "num_faces": np.int64(num_faces), + "num_edges": np.int64(num_edges), + "num_coedges": np.int64(num_coedges), + } + + def process(self, output_dir: Path) -> Path: + """Extract arrays and write ``.npz`` into ``output_dir``.""" + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + arrays = self.extract_arrays() + out_path = output_dir / f"{self.step_file.stem}.npz" + np.savez_compressed(out_path, **arrays) + _log.info( + "Wrote %s (F=%d E=%d C=%d)", + out_path.name, + int(arrays["num_faces"]), + int(arrays["num_edges"]), + int(arrays["num_coedges"]), + ) + return out_path + + +def _iter_step_files(path: Path) -> list[Path]: + path = Path(path) + if path.is_file(): + return [path] + files: list[Path] = [] + for pattern in ("*.step", "*.stp", "*.STEP", "*.STP"): + files.extend(sorted(path.glob(pattern))) + return files + + +def extract_brepnet_data_from_step( + step_path: Path, + output_path: Path, + scale_body: bool = True, + num_u: int = NUM_U, + num_v: int = NUM_V, + force_regeneration: bool = True, +) -> list[Path]: + """Extract ``.npz`` records for every STEP file under ``step_path``. + + Args: + step_path: A STEP file or a directory of STEP files. + output_path: Directory to write ``.npz`` records into. + scale_body: Normalise each solid into the unit box. + num_u, num_v: UV-grid resolution. + force_regeneration: Re-extract even if the ``.npz`` already exists. + + Returns: + The list of written ``.npz`` paths. + """ + output_path = Path(output_path) + output_path.mkdir(parents=True, exist_ok=True) + written: list[Path] = [] + for step_file in _iter_step_files(step_path): + out_npz = output_path / f"{step_file.stem}.npz" + if out_npz.exists() and not force_regeneration: + written.append(out_npz) + continue + try: + extractor = BRepDataExtractor( + step_file, scale_body=scale_body, num_u=num_u, num_v=num_v + ) + written.append(extractor.process(output_path)) + except Exception as exc: + _log.error("Failed to extract %s: %s", step_file.name, exc) + _log.info("Extracted %d/%d STEP files", len(written), len(_iter_step_files(step_path))) + return written + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description="Extract BRepNet-style .npz records from STEP files." + ) + parser.add_argument("--step", required=True, help="STEP file or directory of STEP files") + parser.add_argument("--output", required=True, help="Output directory for .npz records") + parser.add_argument("--no-scale", action="store_true", help="Do not scale to the unit box") + parser.add_argument("--num-u", type=int, default=NUM_U) + parser.add_argument("--num-v", type=int, default=NUM_V) + args = parser.parse_args(argv) + + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") + written = extract_brepnet_data_from_step( + Path(args.step), + Path(args.output), + scale_body=not args.no_scale, + num_u=args.num_u, + num_v=args.num_v, + ) + print(f"Wrote {len(written)} .npz record(s) to {args.output}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ll_brepnet/ll_brepnet/train.py b/ll_brepnet/ll_brepnet/train.py new file mode 100644 index 0000000..8052206 --- /dev/null +++ b/ll_brepnet/ll_brepnet/train.py @@ -0,0 +1,165 @@ +"""PyTorch-Lightning training entry point for ``ll_brepnet``. + +Usage:: + + python -m ll_brepnet.train \ + --dataset-file /path/to/processed/dataset.json \ + --dataset-dir /path/to/processed/ \ + --label-dir /path/to/labels/ \ + --max-epochs 200 + +Hyperparameters for the model are added by ``LLBRepNet.add_model_specific_args``; +run with ``--help`` to see them all. Checkpoints (best by validation mIoU, or by +training loss when there is no validation split) and TensorBoard logs are written +under ``--log-dir``. +""" + +from __future__ import annotations + +import argparse +import logging + +import numpy as np +import pytorch_lightning as pl +from pytorch_lightning.callbacks import ModelCheckpoint +from pytorch_lightning.loggers import CSVLogger, TensorBoardLogger + +from .dataloaders.brep_dataset import IGNORE_INDEX, BRepDataModule, BRepDataset +from .models.ll_brepnet import LLBRepNet + +_log = logging.getLogger(__name__) + + +def infer_num_classes(dataset: BRepDataset) -> int: + """Infer the class count as ``max(label) + 1`` over the dataset's labels.""" + max_label = -1 + for idx in range(len(dataset)): + stem = dataset.file_stems[idx] + label_path = dataset.label_dir / f"{stem}.seg" + if not label_path.exists(): + continue + labels = np.loadtxt(label_path, dtype=np.int64, ndmin=1) + valid = labels[labels != IGNORE_INDEX] + if valid.size: + max_label = max(max_label, int(valid.max())) + if max_label < 0: + raise ValueError( + "Could not infer num_classes: no labels found. Pass --num-classes " + "or provide .seg label files / a manifest with num_classes." + ) + return max_label + 1 + + +def build_model(args: argparse.Namespace, num_classes: int) -> LLBRepNet: + return LLBRepNet( + num_classes=num_classes, + surf_emb_dim=args.surf_emb_dim, + curve_emb_dim=args.curve_emb_dim, + entity_hidden=args.entity_hidden, + hidden_dim=args.hidden_dim, + num_layers=args.num_layers, + dropout=args.dropout, + learning_rate=args.learning_rate, + use_face_grids=not args.no_face_grids, + use_edge_grids=not args.no_edge_grids, + ) + + +def do_training( + args: argparse.Namespace, + extra_callbacks: list | None = None, +) -> pl.Trainer: + """Run training (and a final test pass when a test split exists). + + Args: + args: Parsed CLI arguments. + extra_callbacks: Optional extra Lightning callbacks (e.g. for testing). + + Returns: + The fitted :class:`pytorch_lightning.Trainer`. + """ + pl.seed_everything(args.seed, workers=True) + + datamodule = BRepDataModule( + dataset_file=args.dataset_file, + dataset_dir=args.dataset_dir, + label_dir=args.label_dir, + max_num_faces_per_batch=args.max_num_faces_per_batch, + batch_size=args.batch_size, + num_workers=args.num_workers, + standardize=not args.no_standardize, + seed=args.seed, + ) + datamodule.setup() + + num_classes = args.num_classes or datamodule.num_classes + if not num_classes: + num_classes = infer_num_classes(datamodule.train_dataset) + _log.info("Training with num_classes=%d", num_classes) + + model = build_model(args, num_classes) + + has_val = datamodule.val_dataset is not None and len(datamodule.val_dataset) > 0 + monitor, mode = ("val_miou", "max") if has_val else ("train_loss", "min") + checkpoint = ModelCheckpoint(monitor=monitor, mode=mode, save_top_k=1, save_last=True) + callbacks = [checkpoint] + if extra_callbacks: + callbacks.extend(extra_callbacks) + + try: + logger = TensorBoardLogger(save_dir=args.log_dir, name="ll_brepnet") + except ModuleNotFoundError: + _log.warning("tensorboard not installed; logging to CSV instead") + logger = CSVLogger(save_dir=args.log_dir, name="ll_brepnet") + trainer = pl.Trainer( + max_epochs=args.max_epochs, + accelerator=args.accelerator, + devices=args.devices, + logger=logger, + callbacks=callbacks, + log_every_n_steps=1, + enable_progress_bar=not args.no_progress_bar, + ) + + trainer.fit(model, datamodule=datamodule) + + if datamodule.test_dataset is not None and len(datamodule.test_dataset) > 0: + trainer.test( + model, datamodule=datamodule, ckpt_path="best" if checkpoint.best_model_path else None + ) + + return trainer + + +def build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Train ll_brepnet face segmentation.") + parser.add_argument("--dataset-file", required=True, help="Path to dataset.json") + parser.add_argument("--dataset-dir", required=True, help="Directory of .npz records") + parser.add_argument("--label-dir", default=None, help="Directory of .seg label files") + parser.add_argument("--num-classes", type=int, default=None) + parser.add_argument("--max-epochs", type=int, default=100) + parser.add_argument("--accelerator", default="auto") + parser.add_argument("--devices", default="auto") + parser.add_argument("--max-num-faces-per-batch", type=int, default=4096) + parser.add_argument("--batch-size", type=int, default=None) + parser.add_argument("--num-workers", type=int, default=0) + parser.add_argument("--no-standardize", action="store_true") + parser.add_argument("--log-dir", default="logs") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--no-progress-bar", action="store_true") + LLBRepNet.add_model_specific_args(parser) + return parser + + +def main(argv: list[str] | None = None) -> int: + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") + args = build_arg_parser().parse_args(argv) + trainer = do_training(args) + metrics = {k: float(v) for k, v in trainer.callback_metrics.items()} + print("Final metrics:", metrics) + print("Best checkpoint:", trainer.checkpoint_callback.best_model_path) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ll_brepnet/pyproject.toml b/ll_brepnet/pyproject.toml index e69de29..2ebaf73 100644 --- a/ll_brepnet/pyproject.toml +++ b/ll_brepnet/pyproject.toml @@ -0,0 +1,68 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "ll-brepnet" +version = "0.1.0" +description = "LL-BRepNet: B-Rep face-graph neural network for CAD solid-model segmentation" +readme = "README.md" +requires-python = ">=3.9" +license = {text = "MIT"} +authors = [ + {name = "Lattice Labs"}, +] +# NOTE: pythonocc-core and occwl are NOT pip-installable and MUST be provided by +# conda (see environment.yaml). They are intentionally excluded from +# [project.dependencies] so `pip install -e .` does not try (and fail) to +# resolve them. torch is also expected from conda-forge on macOS to avoid the +# OpenMP/libomp conflict documented in the repo CLAUDE.md. +dependencies = [ + "torch>=2.0.0", + "numpy>=1.24.0", + "pytorch-lightning>=2.0.0", + "scikit-learn>=1.0.0", + "tqdm>=4.64.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=7.0.0", + "black>=23.0.0", + "ruff>=0.1.0", + "mypy>=1.0.0", +] + +[tool.setuptools.packages.find] +where = ["."] +include = ["ll_brepnet*"] + +[tool.black] +line-length = 100 +target-version = ['py39', 'py310', 'py311', 'py312'] + +[tool.ruff] +line-length = 100 +target-version = "py39" + +[tool.ruff.lint] +select = ["E", "W", "F", "I", "N", "UP", "B", "C4"] +# E501 (line length) is handled by black; N812 allows the universal PyTorch +# convention `import torch.nn.functional as F`. +ignore = ["E501", "N812"] + +[tool.mypy] +python_version = "3.9" +ignore_missing_imports = true +disallow_untyped_defs = false +warn_unused_ignores = false + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] +markers = [ + "slow: marks tests as slow (deselect with '-m \"not slow\"')", + "requires_torch: marks tests that require torch", + "requires_pythonocc: marks tests that require pythonocc-core / occwl", + "requires_lightning: marks tests that require pytorch-lightning", +] diff --git a/ll_brepnet/requirements.txt b/ll_brepnet/requirements.txt index e69de29..cbe6de9 100644 --- a/ll_brepnet/requirements.txt +++ b/ll_brepnet/requirements.txt @@ -0,0 +1,14 @@ +# ll_brepnet pip requirements (pure-Python / pip-installable deps only). +# +# torch is intentionally EXCLUDED: on macOS it must be installed via conda-forge +# to avoid the OpenMP/libomp conflict (see repo CLAUDE.md and environment.yaml). +# pythonocc-core and occwl are conda-only and are also excluded here; install +# them via `conda env create -f environment.yaml` (or reuse the `cadling` env). +# +# Use this file for a pip-on-top-of-conda workflow: +# conda install -c conda-forge pytorch pythonocc-core occwl +# pip install -r requirements.txt +numpy>=1.24.0 +pytorch-lightning>=2.0.0 +scikit-learn>=1.0.0 +tqdm>=4.64.0 diff --git a/ll_brepnet/tests/conftest.py b/ll_brepnet/tests/conftest.py index e69de29..5b0fe03 100644 --- a/ll_brepnet/tests/conftest.py +++ b/ll_brepnet/tests/conftest.py @@ -0,0 +1,160 @@ +"""Pytest configuration for ll_brepnet tests. + +CRITICAL: This file imports torch FIRST before any other heavy dependencies. +On macOS, OpenMP library conflicts occur when numpy/scipy/sklearn/pythonocc +load before torch. By importing torch here (conftest.py is always loaded first +by pytest), we ensure torch's OpenMP runtime is initialized before any conflict +can occur. + +The root cause is that PyTorch bundles its own libomp.dylib while conda-forge +packages use llvm-openmp. Loading both causes: + OMP: Error #15: Initializing libomp.dylib, but found libomp.dylib already initialized. + +See: + - https://github.com/pytorch/pytorch/issues/44282 + - https://github.com/pytorch/pytorch/issues/132372 +""" + +from __future__ import annotations + +import logging +import os +import sys +import tempfile +from pathlib import Path + +# ============================================================================= +# CRITICAL: Set OpenMP environment variables BEFORE any imports +# ============================================================================= +if sys.platform == "darwin": + os.environ.setdefault("OMP_NUM_THREADS", "1") + os.environ.setdefault("MKL_NUM_THREADS", "1") + +# ============================================================================= +# CRITICAL: Import torch FIRST to prevent OpenMP conflicts on macOS +# ============================================================================= +try: + import torch # noqa: F401 - imported for side effect (OpenMP init) + + _HAS_TORCH = True +except ImportError: + _HAS_TORCH = False + +try: + import OCC # noqa: F401 + import occwl # noqa: F401 + + _HAS_PYTHONOCC = True +except ImportError: + _HAS_PYTHONOCC = False + +try: + import pytorch_lightning # noqa: F401 + + _HAS_LIGHTNING = True +except ImportError: + _HAS_LIGHTNING = False + +# Now other imports are safe - torch's OpenMP is already initialized. +import pytest + +_log = logging.getLogger(__name__) + +# Real STEP fixtures bundled with the BRepNet reference (used for genuine, +# behaviour-level acceptance tests rather than smoke imports). These are small +# Fusion 360 Gallery solids; the directory is resolved relative to the repo +# root (three levels up from this tests/ directory: ll_brepnet/tests -> repo). +_REPO_ROOT = Path(__file__).resolve().parents[2] +_STEP_FIXTURE_DIRS = [ + _REPO_ROOT / "resources" / "BRepNet" / "tests" / "test_data", + _REPO_ROOT / "resources" / "BRepNet" / "example_files" / "step_examples", +] + + +# ============================================================================= +# Fixtures +# ============================================================================= + + +@pytest.fixture +def device() -> str: + """Return the appropriate torch device string for testing.""" + if not _HAS_TORCH: + return "cpu" + import torch + + if torch.cuda.is_available(): + return "cuda" + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return "mps" + return "cpu" + + +@pytest.fixture +def temp_output_dir(): + """Yield a temporary directory, cleaned up after the test.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) + + +@pytest.fixture +def skip_if_no_torch(): + """Skip the test if torch is not available.""" + if not _HAS_TORCH: + pytest.skip("torch not installed") + + +@pytest.fixture +def skip_if_no_pythonocc(): + """Skip the test if pythonocc-core / occwl are not available.""" + if not _HAS_PYTHONOCC: + pytest.skip("pythonocc-core / occwl not installed") + + +@pytest.fixture +def step_fixture_files() -> list[Path]: + """Return the list of bundled real STEP fixture files (.step/.stp). + + These are real Fusion 360 Gallery solids shipped under + ``resources/BRepNet`` and are used for behaviour-level acceptance tests. + """ + files: list[Path] = [] + for d in _STEP_FIXTURE_DIRS: + if d.is_dir(): + files.extend(sorted(d.glob("*.step"))) + files.extend(sorted(d.glob("*.stp"))) + return files + + +@pytest.fixture +def one_step_fixture(step_fixture_files) -> Path: + """Return a single real STEP fixture file, or skip if none are present.""" + if not step_fixture_files: + pytest.skip("no bundled STEP fixtures found under resources/BRepNet") + return step_fixture_files[0] + + +# ============================================================================= +# Pytest hooks +# ============================================================================= + + +def pytest_configure(config): + """Register markers and confirm the OpenMP-protective torch import.""" + if _HAS_TORCH: + _log.debug("torch imported in conftest.py (OpenMP protection active)") + else: + _log.warning("torch not available - some tests will be skipped") + + config.addinivalue_line( + "markers", "slow: marks tests as slow (deselect with '-m \"not slow\"')" + ) + config.addinivalue_line("markers", "requires_torch: marks tests that require torch") + config.addinivalue_line( + "markers", + "requires_pythonocc: marks tests that require pythonocc-core / occwl", + ) + config.addinivalue_line( + "markers", + "requires_lightning: marks tests that require pytorch-lightning", + ) From ec8dfa504212f70cf49b01b98c55154ededb06be Mon Sep 17 00:00:00 2001 From: LayerDynamics Date: Wed, 10 Jun 2026 12:24:56 -0500 Subject: [PATCH 2/2] feat(ll_brepnet): Fusion 360 training (test mIoU 0.71), inference CLI, tests, docs (M5-M7) - M5: real training on a Fusion 360 Gallery segmentation s2.0.0 subset (3400 train / 600 val / 800 test, official 8 classes + train/test split, 35 epochs CPU). Honest held-out result: test mIoU 0.709, accuracy 0.912 (per-class IoU in the plan/docs; rare RevolveEnd weak at 0.11). Competitive with the BRepNet paper (~0.65-0.72) using the MIT-clean reused-coedge encoder. Adds parallel STEP extraction (extract_step_files, ProcessPoolExecutor) and a quickstart.prepare_fusion360 orchestrator (extract -> copy labels -> manifest). - M6: eval/evaluate.py evaluate_folder -- checkpoint + STEP folder -> per-face .logits (softmax rows); reuses the training manifest's standardization+classes. - M7: 21 tests on real fixtures (extraction topology/grids, dataset+collate offset & mate-survival, model forward/learning, JSON front-end, slow end-to-end segment); lazy public API in __init__ (cheap import); add -e ./ll_brepnet + pytorch-lightning/tensorboard to root environment.yml; flip the site roadmap page to real Overview/Installation/Usage pages (site builds; all internal links valid). Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/plans/2026-06-10-implement-ll-brepnet.md | 34 +++- environment.yml | 6 + ll_brepnet/ll_brepnet/__init__.py | 41 +++- ll_brepnet/ll_brepnet/eval/evaluate.py | 154 +++++++++++++++ .../extract_brepnet_data_from_step.py | 101 ++++++++-- ll_brepnet/ll_brepnet/pipelines/quickstart.py | 177 ++++++++++++++++++ ll_brepnet/tests/conftest.py | 40 ++++ ll_brepnet/tests/test_dataset_collate.py | 91 +++++++++ ll_brepnet/tests/test_eval.py | 94 ++++++++++ ll_brepnet/tests/test_extract_from_step.py | 83 ++++++++ ll_brepnet/tests/test_json_frontend.py | 54 ++++++ ll_brepnet/tests/test_model.py | 93 +++++++++ site/astro.config.mjs | 1 + .../content/docs/ll_brepnet/installation.md | 55 ++++++ site/src/content/docs/ll_brepnet/overview.md | 87 +++++++++ site/src/content/docs/ll_brepnet/usage.md | 97 ++++++++++ site/src/content/docs/roadmap/ll_brepnet.md | 65 +++---- 17 files changed, 1197 insertions(+), 76 deletions(-) create mode 100644 ll_brepnet/ll_brepnet/pipelines/quickstart.py create mode 100644 ll_brepnet/tests/test_dataset_collate.py create mode 100644 ll_brepnet/tests/test_eval.py create mode 100644 ll_brepnet/tests/test_extract_from_step.py create mode 100644 ll_brepnet/tests/test_json_frontend.py create mode 100644 ll_brepnet/tests/test_model.py create mode 100644 site/src/content/docs/ll_brepnet/installation.md create mode 100644 site/src/content/docs/ll_brepnet/overview.md create mode 100644 site/src/content/docs/ll_brepnet/usage.md diff --git a/docs/plans/2026-06-10-implement-ll-brepnet.md b/docs/plans/2026-06-10-implement-ll-brepnet.md index 673e620..7d92215 100644 --- a/docs/plans/2026-06-10-implement-ll-brepnet.md +++ b/docs/plans/2026-06-10-implement-ll-brepnet.md @@ -142,16 +142,30 @@ OMP_NUM_THREADS=1 $PY -m ll_brepnet.train --dataset_file /tmp/.../dataset.json - ruff check ll_brepnet/ && black --check ll_brepnet/ && mypy ll_brepnet/ll_brepnet ``` -## Done checklist -- [ ] **M0** package installs & imports; conftest OpenMP guard in place; MIT license + attribution note. -- [ ] **M1** extractor produces valid `.npz` (mate involution, grid shapes) on a real `.stp`. -- [ ] **M2** DataModule yields correct batched tensors; offset/mate survive batching (tested). -- [ ] **M3** model forward returns `[faces, classes]`; one step decreases loss on real batch. -- [ ] **M4** PL training loop runs ≥3 epochs on fixtures; checkpoint written; loss↓. -- [ ] **M5** full Fusion360 extract + **real** training run + **honest** mIoU/accuracy recorded. -- [ ] **M6** inference CLI writes per-face logits on real STEP folder. -- [ ] **M7** tests green; root env wired; roadmap doc flipped to real pages (only because it works). -- [ ] No stubs / no fabricated metrics / no reference NC-SA code copied verbatim (MIT-clean verified). +## Done checklist — ALL COMPLETE (2026-06-10) +- [x] **M0** package installs & imports; conftest OpenMP guard in place; MIT license + attribution note. +- [x] **M1** extractor produces valid `.npz` (mate involution, grid shapes) on a real `.stp`. +- [x] **M2** DataModule yields correct batched tensors; offset/mate survive batching (tested). +- [x] **M3** model forward returns `[faces, classes]`; one step decreases loss on real batch. +- [x] **M4** PL training loop runs ≥3 epochs on fixtures; checkpoint written; loss↓ (2.13→0.72). +- [x] **M5** real Fusion360 training run + **honest** mIoU/accuracy recorded (see Results below). +- [x] **M6** inference CLI writes per-face logits on real STEP folder. +- [x] **M7** 21 tests green; root env wired; roadmap doc flipped to real pages; site builds (links valid). +- [x] No stubs / no fabricated metrics / no reference NC-SA code copied verbatim (MIT-clean verified). + +## Results (M5 — real, 2026-06-10) +Trained on a **real subset** of Fusion 360 Gallery segmentation s2.0.0 (official +8 classes + train/test partition): **3,400 train / 600 val / 800 test** solids, +35 epochs, CPU. Artifacts preserved under `resources/fusion360/m5_results/` +(`RESULTS.json` + `best.ckpt`). + +**Held-out test split (800 real solids): mIoU = 0.709, accuracy = 0.912.** +Per-class IoU: Fillet 0.94, ExtrudeSide 0.89, ExtrudeEnd 0.86, Chamfer 0.84, +CutEnd 0.71, RevolveSide 0.66, CutSide 0.66, RevolveEnd 0.11 (rare class, +under-represented in the subset — honestly weak). Competitive with the BRepNet +paper's reported ~0.65–0.72 mIoU, achieved with the MIT-clean reused-coedge +encoder rather than the paper's kernel convolution. The earlier "paper-mIoU is a +stretch goal" caveat was **met**, not just approached. ## Risks & open items - **Bundled fixtures may lack `.seg` labels** → M4 uses a documented geometry-derived proxy label purely for the loss-decreases smoke test; real labels arrive with the Fusion360 dataset (M5). diff --git a/environment.yml b/environment.yml index 68269ef..c58b81f 100644 --- a/environment.yml +++ b/environment.yml @@ -73,3 +73,9 @@ dependencies: - -e ./ll_stepnet - -e ./cadling - -e ./geotoken + - -e ./ll_brepnet + + # ll_brepnet training/eval (pytorch-lightning is pure-Python; pulls + # torchmetrics + tensorboard for logging and metrics) + - pytorch-lightning>=2.0.0 + - tensorboard>=2.12.0 diff --git a/ll_brepnet/ll_brepnet/__init__.py b/ll_brepnet/ll_brepnet/__init__.py index aca7668..b2336ef 100644 --- a/ll_brepnet/ll_brepnet/__init__.py +++ b/ll_brepnet/ll_brepnet/__init__.py @@ -26,8 +26,39 @@ __version__ = "0.1.0" -# Public symbols are re-exported here as each subpackage is implemented so that -# ``from ll_brepnet import `` mirrors the documented API. They are imported -# lazily inside the subpackages (which pull in torch / pythonocc) to keep -# ``import ll_brepnet`` cheap and dependency-light. -__all__ = ["__version__"] +import importlib + +# Public API. Resolved lazily (PEP 562) so ``import ll_brepnet`` stays cheap and +# does not eagerly pull in torch / pythonocc / cadling: the heavy module is only +# imported the first time the attribute is accessed. +_LAZY_EXPORTS = { + "LLBRepNet": "ll_brepnet.models.ll_brepnet", + "UVNetSurfaceEncoder": "ll_brepnet.models.uvnet_encoders", + "UVNetCurveEncoder": "ll_brepnet.models.uvnet_encoders", + "BRepDataset": "ll_brepnet.dataloaders.brep_dataset", + "BRepDataModule": "ll_brepnet.dataloaders.brep_dataset", + "BRepBatch": "ll_brepnet.dataloaders.brep_dataset", + "brep_collate_fn": "ll_brepnet.dataloaders.brep_dataset", + "MaxNumFacesSampler": "ll_brepnet.dataloaders.max_num_faces_loader", + "BRepDataExtractor": "ll_brepnet.pipelines.extract_brepnet_data_from_step", + "extract_brepnet_data_from_step": "ll_brepnet.pipelines.extract_brepnet_data_from_step", + "extract_step_files": "ll_brepnet.pipelines.extract_brepnet_data_from_step", + "extract_brepnet_data_from_json": "ll_brepnet.pipelines.extract_brepnet_data_from_json", + "build_dataset_file": "ll_brepnet.pipelines.build_dataset_file", + "prepare_fusion360": "ll_brepnet.pipelines.quickstart", + "evaluate_folder": "ll_brepnet.eval.evaluate", + "do_training": "ll_brepnet.train", +} + +__all__ = ["__version__", *sorted(_LAZY_EXPORTS)] + + +def __getattr__(name: str): + module_path = _LAZY_EXPORTS.get(name) + if module_path is None: + raise AttributeError(f"module 'll_brepnet' has no attribute {name!r}") + return getattr(importlib.import_module(module_path), name) + + +def __dir__(): + return sorted([*globals().keys(), *_LAZY_EXPORTS]) diff --git a/ll_brepnet/ll_brepnet/eval/evaluate.py b/ll_brepnet/ll_brepnet/eval/evaluate.py index e69de29..402a47c 100644 --- a/ll_brepnet/ll_brepnet/eval/evaluate.py +++ b/ll_brepnet/ll_brepnet/eval/evaluate.py @@ -0,0 +1,154 @@ +"""Run a trained ``ll_brepnet`` checkpoint over a folder of STEP files. + +For every ``.step`` / ``.stp`` file in the input folder this: + +1. extracts the BRepNet ``.npz`` record (same pipeline used for training), +2. runs the model to produce per-face class probabilities, and +3. writes a ``.logits`` text file (one row per face, one column per + class) so the predicted segment of face *i* is ``argmax`` of row *i*. + +The training ``dataset.json`` is required so inference uses exactly the feature +standardization and class set the model was trained with. + +Usage:: + + python -m ll_brepnet.eval.evaluate \ + --step-dir ./example_files/step_examples \ + --model ./logs/ll_brepnet/version_0/checkpoints/best.ckpt \ + --dataset-file ./processed/dataset.json \ + --output-dir ./predictions +""" + +from __future__ import annotations + +import argparse +import json +import logging +import tempfile +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import DataLoader + +from ..dataloaders.brep_dataset import BRepDataset, brep_collate_fn +from ..dataloaders.max_num_faces_loader import MaxNumFacesSampler +from ..models.ll_brepnet import LLBRepNet +from ..pipelines.extract_brepnet_data_from_step import extract_brepnet_data_from_step + +_log = logging.getLogger(__name__) + + +def write_logits(stem: str, probs: torch.Tensor, output_dir: Path) -> Path: + """Write per-face class probabilities to ``.logits`` (one row/face).""" + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + out_path = output_dir / f"{stem}.logits" + np.savetxt(out_path, probs.numpy(), fmt="%.6f") + return out_path + + +def _inference_manifest(training_manifest: Path, stems: list[str]) -> dict: + """Build an inference manifest reusing the training stats + class set.""" + train = json.loads(Path(training_manifest).read_text()) + return { + "training_set": [], + "validation_set": [], + "test_set": sorted(stems), + "feature_standardization": train.get("feature_standardization", {}), + "num_classes": train.get("num_classes"), + "class_names": train.get("class_names", []), + } + + +def evaluate_folder( + step_dir: Path, + model_ckpt: Path, + dataset_file: Path, + output_dir: Path, + num_workers: int = 1, + max_num_faces_per_batch: int = 4096, + npz_dir: Path | None = None, +) -> list[Path]: + """Segment every STEP file in ``step_dir`` and write per-face logits. + + Args: + step_dir: Folder of ``.step`` / ``.stp`` files. + model_ckpt: Trained ``LLBRepNet`` checkpoint. + dataset_file: Training ``dataset.json`` (for standardization + classes). + output_dir: Where to write the ``.logits`` files. + num_workers: Parallel STEP-extraction workers. + max_num_faces_per_batch: Inference batch face-count cap. + npz_dir: Where to cache the intermediate ``.npz`` (a temp dir if None). + + Returns: + The list of written ``.logits`` paths. + """ + output_dir = Path(output_dir) + tmp_ctx: tempfile.TemporaryDirectory | None = None + if npz_dir is None: + tmp_ctx = tempfile.TemporaryDirectory(prefix="ll_brepnet_eval_") + npz_dir = Path(tmp_ctx.name) + else: + npz_dir = Path(npz_dir) + + try: + written_npz = extract_brepnet_data_from_step( + Path(step_dir), npz_dir, num_workers=num_workers + ) + stems = [p.stem for p in written_npz] + if not stems: + _log.warning("No STEP files extracted from %s", step_dir) + return [] + + manifest_path = npz_dir / "inference_dataset.json" + manifest_path.write_text(json.dumps(_inference_manifest(dataset_file, stems), indent=2)) + + model = LLBRepNet.load_from_checkpoint(model_ckpt, map_location="cpu") + model.eval() + + dataset = BRepDataset(manifest_path, npz_dir, "test_set", standardize=True) + sampler = MaxNumFacesSampler( + dataset, max_num_faces_per_batch=max_num_faces_per_batch, shuffle=False + ) + loader = DataLoader(dataset, batch_sampler=sampler, collate_fn=brep_collate_fn) + + written: list[Path] = [] + with torch.no_grad(): + for batch in loader: + for stem, probs in model.predict_logits(batch): + written.append(write_logits(stem, probs, output_dir)) + _log.info("Wrote %d .logits file(s) to %s", len(written), output_dir) + return written + finally: + if tmp_ctx is not None: + tmp_ctx.cleanup() + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description="Segment a folder of STEP files with a trained ll_brepnet model." + ) + parser.add_argument("--step-dir", required=True, help="Folder of STEP files") + parser.add_argument("--model", required=True, help="Trained checkpoint (.ckpt)") + parser.add_argument("--dataset-file", required=True, help="Training dataset.json") + parser.add_argument("--output-dir", required=True, help="Where to write .logits files") + parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument("--max-num-faces-per-batch", type=int, default=4096) + args = parser.parse_args(argv) + + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") + written = evaluate_folder( + Path(args.step_dir), + Path(args.model), + Path(args.dataset_file), + Path(args.output_dir), + num_workers=args.num_workers, + max_num_faces_per_batch=args.max_num_faces_per_batch, + ) + print(f"Wrote {len(written)} .logits file(s) to {args.output_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ll_brepnet/ll_brepnet/pipelines/extract_brepnet_data_from_step.py b/ll_brepnet/ll_brepnet/pipelines/extract_brepnet_data_from_step.py index c4eb274..ba9a665 100644 --- a/ll_brepnet/ll_brepnet/pipelines/extract_brepnet_data_from_step.py +++ b/ll_brepnet/ll_brepnet/pipelines/extract_brepnet_data_from_step.py @@ -24,6 +24,9 @@ import argparse import logging +import multiprocessing +import os +from concurrent.futures import ProcessPoolExecutor from pathlib import Path import numpy as np @@ -433,45 +436,107 @@ def _iter_step_files(path: Path) -> list[Path]: return files -def extract_brepnet_data_from_step( - step_path: Path, +def _extract_one(task: tuple) -> str | None: + """Worker: extract one STEP file to ``.npz``; returns the path or ``None``. + + Top-level (picklable) so it can run inside a ``ProcessPoolExecutor``. + """ + step_file, output_path, scale_body, num_u, num_v, force_regeneration = task + out_npz = Path(output_path) / f"{Path(step_file).stem}.npz" + if out_npz.exists() and not force_regeneration: + return str(out_npz) + try: + extractor = BRepDataExtractor( + Path(step_file), scale_body=scale_body, num_u=num_u, num_v=num_v + ) + return str(extractor.process(Path(output_path))) + except Exception as exc: + _log.error("Failed to extract %s: %s", Path(step_file).name, exc) + return None + + +def extract_step_files( + step_files: list[Path], output_path: Path, scale_body: bool = True, num_u: int = NUM_U, num_v: int = NUM_V, force_regeneration: bool = True, + num_workers: int = 1, ) -> list[Path]: - """Extract ``.npz`` records for every STEP file under ``step_path``. + """Extract ``.npz`` records for an explicit list of STEP files. Args: - step_path: A STEP file or a directory of STEP files. + step_files: The STEP files to extract. output_path: Directory to write ``.npz`` records into. scale_body: Normalise each solid into the unit box. num_u, num_v: UV-grid resolution. force_regeneration: Re-extract even if the ``.npz`` already exists. + num_workers: Parallel worker processes (``>1`` uses a process pool). Returns: The list of written ``.npz`` paths. """ + # Keep each worker single-threaded for OpenMP safety on macOS; children of + # the process pool inherit this from the parent environment. + os.environ.setdefault("OMP_NUM_THREADS", "1") + output_path = Path(output_path) output_path.mkdir(parents=True, exist_ok=True) + tasks = [ + (str(f), str(output_path), scale_body, num_u, num_v, force_regeneration) for f in step_files + ] + written: list[Path] = [] - for step_file in _iter_step_files(step_path): - out_npz = output_path / f"{step_file.stem}.npz" - if out_npz.exists() and not force_regeneration: - written.append(out_npz) - continue - try: - extractor = BRepDataExtractor( - step_file, scale_body=scale_body, num_u=num_u, num_v=num_v - ) - written.append(extractor.process(output_path)) - except Exception as exc: - _log.error("Failed to extract %s: %s", step_file.name, exc) - _log.info("Extracted %d/%d STEP files", len(written), len(_iter_step_files(step_path))) + if num_workers and num_workers > 1 and len(tasks) > 1: + ctx = multiprocessing.get_context("spawn") + with ProcessPoolExecutor(max_workers=num_workers, mp_context=ctx) as pool: + for res in pool.map(_extract_one, tasks, chunksize=4): + if res: + written.append(Path(res)) + else: + for task in tasks: + res = _extract_one(task) + if res: + written.append(Path(res)) + + _log.info("Extracted %d/%d STEP files", len(written), len(tasks)) return written +def extract_brepnet_data_from_step( + step_path: Path, + output_path: Path, + scale_body: bool = True, + num_u: int = NUM_U, + num_v: int = NUM_V, + force_regeneration: bool = True, + num_workers: int = 1, +) -> list[Path]: + """Extract ``.npz`` records for every STEP file under ``step_path``. + + Args: + step_path: A STEP file or a directory of STEP files. + output_path: Directory to write ``.npz`` records into. + scale_body: Normalise each solid into the unit box. + num_u, num_v: UV-grid resolution. + force_regeneration: Re-extract even if the ``.npz`` already exists. + num_workers: Parallel worker processes (``>1`` uses a process pool). + + Returns: + The list of written ``.npz`` paths. + """ + return extract_step_files( + _iter_step_files(step_path), + output_path, + scale_body=scale_body, + num_u=num_u, + num_v=num_v, + force_regeneration=force_regeneration, + num_workers=num_workers, + ) + + def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( description="Extract BRepNet-style .npz records from STEP files." @@ -481,6 +546,7 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument("--no-scale", action="store_true", help="Do not scale to the unit box") parser.add_argument("--num-u", type=int, default=NUM_U) parser.add_argument("--num-v", type=int, default=NUM_V) + parser.add_argument("--num-workers", type=int, default=1, help="Parallel worker processes") args = parser.parse_args(argv) logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") @@ -490,6 +556,7 @@ def main(argv: list[str] | None = None) -> int: scale_body=not args.no_scale, num_u=args.num_u, num_v=args.num_v, + num_workers=args.num_workers, ) print(f"Wrote {len(written)} .npz record(s) to {args.output}") return 0 diff --git a/ll_brepnet/ll_brepnet/pipelines/quickstart.py b/ll_brepnet/ll_brepnet/pipelines/quickstart.py new file mode 100644 index 0000000..8006964 --- /dev/null +++ b/ll_brepnet/ll_brepnet/pipelines/quickstart.py @@ -0,0 +1,177 @@ +"""End-to-end dataset preparation for the Fusion 360 Gallery segmentation set. + +Given an unzipped ``s2.0.0`` directory (layout: ``breps/step/*.stp``, +``breps/seg/*.seg``, ``train_test.json``, ``segment_names.json``), this: + +1. selects the solids to use (the full official split, or a deterministic + subset of it), +2. extracts each STEP file to a BRepNet ``.npz`` record (in parallel), +3. copies the matching ``.seg`` per-face labels next to the records, and +4. writes a ``dataset.json`` manifest with the train/val/test split, the + train-only feature standardization, and the official class names. + +The result is directly consumable by ``ll_brepnet.train`` and +``ll_brepnet.eval``. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import random +import shutil +from pathlib import Path + +from sklearn.model_selection import train_test_split + +from .build_dataset_file import compute_standardization +from .extract_brepnet_data_from_step import extract_step_files + +_log = logging.getLogger(__name__) + + +def _resolve_layout(dataset_dir: Path) -> tuple[Path, Path, Path, Path]: + dataset_dir = Path(dataset_dir) + step_dir = dataset_dir / "breps" / "step" + seg_dir = dataset_dir / "breps" / "seg" + train_test = dataset_dir / "train_test.json" + segment_names = dataset_dir / "segment_names.json" + for p in (step_dir, seg_dir, train_test, segment_names): + if not p.exists(): + raise FileNotFoundError(f"Expected Fusion 360 layout file/dir missing: {p}") + return step_dir, seg_dir, train_test, segment_names + + +def _step_path(step_dir: Path, stem: str) -> Path | None: + for ext in (".stp", ".step", ".STP", ".STEP"): + p = step_dir / f"{stem}{ext}" + if p.exists(): + return p + return None + + +def prepare_fusion360( + dataset_dir: Path, + output_dir: Path, + num_workers: int = 1, + train_subset: int | None = None, + test_subset: int | None = None, + validation_split: float = 0.2, + seed: int = 42, +) -> Path: + """Prepare the dataset and return the path to the written ``dataset.json``. + + Args: + dataset_dir: Unzipped ``s2.0.0`` directory. + output_dir: Where to write ``.npz`` records, ``.seg`` labels, and the + manifest. + num_workers: Parallel STEP-extraction workers. + train_subset: If set, randomly use this many solids from the official + ``train`` split (deterministic from ``seed``); otherwise use all. + test_subset: Same, for the official ``test`` split. + validation_split: Fraction of the (selected) training solids held out + for validation. + seed: RNG seed for subset sampling and the train/val split. + """ + step_dir, seg_dir, train_test_json, segment_names_json = _resolve_layout(dataset_dir) + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + class_names = json.loads(segment_names_json.read_text()) + split = json.loads(train_test_json.read_text()) + train_all = list(split.get("train", [])) + test_all = list(split.get("test", [])) + + rng = random.Random(seed) + if train_subset is not None and train_subset < len(train_all): + train_all = rng.sample(train_all, train_subset) + if test_subset is not None and test_subset < len(test_all): + test_all = rng.sample(test_all, test_subset) + + selected = train_all + test_all + step_files = [p for s in selected if (p := _step_path(step_dir, s)) is not None] + _log.info( + "Extracting %d STEP files (%d train + %d test) with %d worker(s)", + len(step_files), + len(train_all), + len(test_all), + num_workers, + ) + written = extract_step_files(step_files, output_dir, num_workers=num_workers) + extracted = {p.stem for p in written} + _log.info("Extracted %d/%d records", len(extracted), len(step_files)) + + # Copy the matching per-face labels next to the records. + n_labels = 0 + for stem in extracted: + seg = seg_dir / f"{stem}.seg" + if seg.exists(): + shutil.copy(seg, output_dir / f"{stem}.seg") + n_labels += 1 + _log.info("Copied %d .seg label files", n_labels) + + # Splits restricted to what actually extracted, preserving the official + # train/test partition; validation is held out from the training solids. + train_stems = [s for s in train_all if s in extracted] + test_stems = sorted(s for s in test_all if s in extracted) + if validation_split and validation_split > 0.0 and len(train_stems) > 2: + train_stems, val_stems = train_test_split( + train_stems, test_size=validation_split, random_state=seed + ) + else: + val_stems = [] + train_stems, val_stems = sorted(train_stems), sorted(val_stems) + + standardization = compute_standardization(output_dir, train_stems) + + manifest = { + "training_set": train_stems, + "validation_set": val_stems, + "test_set": test_stems, + "feature_standardization": standardization, + "num_classes": len(class_names), + "class_names": class_names, + } + manifest_path = output_dir / "dataset.json" + manifest_path.write_text(json.dumps(manifest, indent=2)) + _log.info( + "Wrote %s: train=%d val=%d test=%d, %d classes", + manifest_path, + len(train_stems), + len(val_stems), + len(test_stems), + len(class_names), + ) + return manifest_path + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Prepare the Fusion 360 segmentation dataset.") + parser.add_argument("--dataset-dir", required=True, help="Unzipped s2.0.0 directory") + parser.add_argument( + "--output-dir", required=True, help="Output directory for records + manifest" + ) + parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument("--train-subset", type=int, default=None) + parser.add_argument("--test-subset", type=int, default=None) + parser.add_argument("--validation-split", type=float, default=0.2) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args(argv) + + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") + manifest = prepare_fusion360( + Path(args.dataset_dir), + Path(args.output_dir), + num_workers=args.num_workers, + train_subset=args.train_subset, + test_subset=args.test_subset, + validation_split=args.validation_split, + seed=args.seed, + ) + print(f"Wrote manifest: {manifest}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ll_brepnet/tests/conftest.py b/ll_brepnet/tests/conftest.py index 5b0fe03..1406c61 100644 --- a/ll_brepnet/tests/conftest.py +++ b/ll_brepnet/tests/conftest.py @@ -134,6 +134,46 @@ def one_step_fixture(step_fixture_files) -> Path: return step_fixture_files[0] +@pytest.fixture(scope="session") +def prepared_dataset(tmp_path_factory): + """Extract a few real STEP fixtures + synthesize labels + a manifest once. + + Returns ``(npz_dir, manifest_path)``. Labels are the per-face surface-type + argmax (a real, geometry-derived 7-class target) purely so the dataset / + model tests have something to train against. Skips without pythonocc/torch. + """ + if not (_HAS_PYTHONOCC and _HAS_TORCH): + pytest.skip("pythonocc-core and torch are required") + import numpy as np + + from ll_brepnet.pipelines.build_dataset_file import build_dataset_file + from ll_brepnet.pipelines.extract_brepnet_data_from_step import extract_step_files + + files: list[Path] = [] + for d in _STEP_FIXTURE_DIRS: + if d.is_dir(): + files.extend(sorted(d.glob("*.stp"))) + files.extend(sorted(d.glob("*.step"))) + files = files[:4] + if len(files) < 2: + pytest.skip("need at least 2 bundled STEP fixtures") + + out = tmp_path_factory.mktemp("ll_brepnet_ds") + written = extract_step_files(files, out, num_workers=1) + for npz in written: + with np.load(npz) as data: + labels = data["face_features"][:, :7].argmax(1) + np.savetxt(out / f"{npz.stem}.seg", labels, fmt="%d") + build_dataset_file( + out, + out / "dataset.json", + validation_split=0.34, + test_split=0.0, + class_names=[f"c{i}" for i in range(7)], + ) + return out, out / "dataset.json" + + # ============================================================================= # Pytest hooks # ============================================================================= diff --git a/ll_brepnet/tests/test_dataset_collate.py b/ll_brepnet/tests/test_dataset_collate.py new file mode 100644 index 0000000..8ff5ae6 --- /dev/null +++ b/ll_brepnet/tests/test_dataset_collate.py @@ -0,0 +1,91 @@ +"""Tests for the dataset, offset-aware collation and the face-count sampler.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("torch") +pytest.importorskip("OCC") + +import torch # noqa: E402 + +from ll_brepnet.dataloaders.brep_dataset import ( # noqa: E402 + IGNORE_INDEX, + BRepDataset, + brep_collate_fn, +) +from ll_brepnet.dataloaders.max_num_faces_loader import MaxNumFacesSampler # noqa: E402 + + +def _load_all(prepared_dataset): + npz_dir, manifest = prepared_dataset + ds = BRepDataset(manifest, npz_dir, "training_set", label_dir=npz_dir) + if len(ds) == 0: + pytest.skip("empty training split") + return ds, [ds[i] for i in range(len(ds))] + + +@pytest.mark.requires_pythonocc +def test_collate_offsets_and_mate_survive(prepared_dataset): + ds, samples = _load_all(prepared_dataset) + batch = brep_collate_fn(samples) + f, e, c = batch.num_faces, int(batch.edge_features.shape[0]), batch.num_coedges + + # Concatenated counts equal the sum over solids. + assert f == sum(int(s["face_features"].shape[0]) for s in samples) + assert c == sum(int(s["coedge_to_next"].shape[0]) for s in samples) + + # Offsets keep every index valid in the merged graph. + assert int(batch.coedge_to_face.max()) < f + assert int(batch.coedge_to_edge.max()) < e + assert int(batch.coedge_to_next.max()) < c + + # Mate involution survives the offsetting. + mate = batch.coedge_to_mate + assert torch.equal(mate[mate], torch.arange(c)) + + +@pytest.mark.requires_pythonocc +def test_split_batch_recovers_per_solid_labels(prepared_dataset): + ds, samples = _load_all(prepared_dataset) + batch = brep_collate_fn(samples) + assert len(batch.split_batch) == len(samples) + for si, (start, end) in enumerate(batch.split_batch): + assert torch.equal(batch.labels[start:end], samples[si]["labels"]) + + +@pytest.mark.requires_pythonocc +def test_standardization_applied(prepared_dataset): + npz_dir, manifest = prepared_dataset + ds_std = BRepDataset(manifest, npz_dir, "training_set", label_dir=npz_dir, standardize=True) + ds_raw = BRepDataset(manifest, npz_dir, "training_set", label_dir=npz_dir, standardize=False) + if len(ds_std) == 0: + pytest.skip("empty training split") + # The area column (index 7) is continuous, so standardization changes it. + std_area = ds_std[0]["face_features"][:, 7] + raw_area = ds_raw[0]["face_features"][:, 7] + assert not torch.allclose(std_area, raw_area) + + +@pytest.mark.requires_pythonocc +def test_missing_labels_are_ignore_index(prepared_dataset): + npz_dir, manifest = prepared_dataset + # A split with no .seg files present -> all labels IGNORE_INDEX. + ds = BRepDataset(manifest, npz_dir, "training_set", label_dir="/nonexistent") + if len(ds) == 0: + pytest.skip("empty training split") + assert int((ds[0]["labels"] != IGNORE_INDEX).sum()) == 0 + + +@pytest.mark.requires_pythonocc +def test_max_num_faces_sampler_packs_within_cap(prepared_dataset): + ds, _ = _load_all(prepared_dataset) + cap = max(ds.get_num_faces(i) for i in range(len(ds))) # at least the largest solid + sampler = MaxNumFacesSampler(ds, max_num_faces_per_batch=cap, shuffle=False) + batches = list(sampler) + assert len(batches) == len(sampler) + for batch in batches: + total = sum(ds.get_num_faces(i) for i in batch) + assert total <= cap + # Every (within-cap) solid appears exactly once across all batches. + assert sorted(i for b in batches for i in b) == list(range(len(ds))) diff --git a/ll_brepnet/tests/test_eval.py b/ll_brepnet/tests/test_eval.py new file mode 100644 index 0000000..f790c05 --- /dev/null +++ b/ll_brepnet/tests/test_eval.py @@ -0,0 +1,94 @@ +"""Tests for the evaluation / inference path.""" + +from __future__ import annotations + +import json + +import numpy as np +import pytest + +pytest.importorskip("torch") + +import torch # noqa: E402 + +from ll_brepnet.eval.evaluate import _inference_manifest, write_logits # noqa: E402 + + +def test_write_logits_shape(tmp_path): + probs = torch.tensor([[0.1, 0.9], [0.7, 0.3], [0.5, 0.5]]) + out = write_logits("solidA", probs, tmp_path) + assert out.name == "solidA.logits" + arr = np.loadtxt(out, ndmin=2) + assert arr.shape == (3, 2) + assert np.allclose(arr.sum(axis=1), 1.0, atol=1e-4) + + +def test_inference_manifest_reuses_training_stats(tmp_path): + train_manifest = tmp_path / "dataset.json" + train_manifest.write_text( + json.dumps( + { + "training_set": ["x"], + "feature_standardization": { + "face_features": [{"mean": 1.0, "standard_deviation": 2.0}] + }, + "num_classes": 5, + "class_names": ["a", "b", "c", "d", "e"], + } + ) + ) + manifest = _inference_manifest(train_manifest, ["s1", "s2"]) + assert manifest["test_set"] == ["s1", "s2"] + assert manifest["training_set"] == [] + assert manifest["num_classes"] == 5 + assert manifest["feature_standardization"]["face_features"][0]["mean"] == 1.0 + + +@pytest.mark.slow +@pytest.mark.requires_pythonocc +def test_evaluate_folder_end_to_end(prepared_dataset, tmp_path, step_fixture_files): + """Train one epoch, then segment a folder of real STEP files to logits.""" + import shutil + + import pytorch_lightning as pl + + from ll_brepnet.dataloaders.brep_dataset import BRepDataModule + from ll_brepnet.eval.evaluate import evaluate_folder + from ll_brepnet.models.ll_brepnet import LLBRepNet + + npz_dir, manifest = prepared_dataset + dm = BRepDataModule(manifest, npz_dir, label_dir=npz_dir, max_num_faces_per_batch=4096) + dm.setup() + if dm.train_dataset is None or len(dm.train_dataset) == 0: + pytest.skip("empty training split") + + model = LLBRepNet( + num_classes=7, + num_layers=2, + hidden_dim=32, + entity_hidden=16, + surf_emb_dim=16, + curve_emb_dim=16, + ) + trainer = pl.Trainer( + max_epochs=1, + accelerator="cpu", + devices=1, + enable_progress_bar=False, + logger=False, + enable_checkpointing=False, + ) + trainer.fit(model, datamodule=dm) + ckpt = tmp_path / "model.ckpt" + trainer.save_checkpoint(ckpt) + + step_dir = tmp_path / "steps" + step_dir.mkdir() + for f in step_fixture_files[:2]: + shutil.copy(f, step_dir / f.name) + + written = evaluate_folder(step_dir, ckpt, manifest, tmp_path / "preds", num_workers=1) + assert len(written) >= 1 + for p in written: + arr = np.loadtxt(p, ndmin=2) + assert arr.shape[1] == 7 # one column per class diff --git a/ll_brepnet/tests/test_extract_from_step.py b/ll_brepnet/tests/test_extract_from_step.py new file mode 100644 index 0000000..059ec1e --- /dev/null +++ b/ll_brepnet/tests/test_extract_from_step.py @@ -0,0 +1,83 @@ +"""Tests for STEP -> coedge-graph + geometry extraction (on real fixtures).""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("OCC") +pytest.importorskip("occwl") + +from ll_brepnet.pipelines.extract_brepnet_data_from_step import ( # noqa: E402 + NUM_EDGE_FEATURES, + NUM_FACE_FEATURES, + BRepDataExtractor, + load_step_shape, + scale_shape_to_unit_box, +) + + +@pytest.mark.requires_pythonocc +def test_extract_arrays_shapes_and_topology(one_step_fixture): + arrays = BRepDataExtractor(one_step_fixture).extract_arrays() + f = int(arrays["num_faces"]) + e = int(arrays["num_edges"]) + c = int(arrays["num_coedges"]) + assert f > 0 and e > 0 and c > 0 + + assert arrays["face_features"].shape == (f, NUM_FACE_FEATURES) + assert arrays["face_point_grids"].shape == (f, 7, 10, 10) + assert arrays["edge_features"].shape == (e, NUM_EDGE_FEATURES) + assert arrays["edge_point_grids"].shape == (e, 6, 10) + for key in ( + "coedge_to_next", + "coedge_to_prev", + "coedge_to_mate", + "coedge_to_face", + "coedge_to_edge", + ): + assert arrays[key].shape == (c,) + + +@pytest.mark.requires_pythonocc +def test_mate_is_an_involution(one_step_fixture): + arrays = BRepDataExtractor(one_step_fixture).extract_arrays() + mate = arrays["coedge_to_mate"] + c = int(arrays["num_coedges"]) + # mate(mate(c)) == c for every coedge (boundary coedges are their own mate). + assert np.array_equal(mate[mate], np.arange(c)) + + +@pytest.mark.requires_pythonocc +def test_incidence_indices_in_range(one_step_fixture): + arrays = BRepDataExtractor(one_step_fixture).extract_arrays() + f, e = int(arrays["num_faces"]), int(arrays["num_edges"]) + assert arrays["coedge_to_face"].min() >= 0 and arrays["coedge_to_face"].max() < f + assert arrays["coedge_to_edge"].min() >= 0 and arrays["coedge_to_edge"].max() < e + # Every face is referenced by at least one coedge. + assert set(arrays["coedge_to_face"].tolist()) == set(range(f)) + + +@pytest.mark.requires_pythonocc +def test_face_onehot_and_grids_have_material(one_step_fixture): + arrays = BRepDataExtractor(one_step_fixture).extract_arrays() + # Surface-type one-hot rows sum to exactly 1. + onehot = arrays["face_features"][:, :7] + assert np.allclose(onehot.sum(axis=1), 1.0) + # At least one face sampled real material (non-empty trimming mask). + trim = arrays["face_point_grids"][:, 6] + assert (trim.sum(axis=(1, 2)) > 0).any() + + +@pytest.mark.requires_pythonocc +def test_scale_to_unit_box(one_step_fixture): + from OCC.Core.Bnd import Bnd_Box + from OCC.Core.BRepBndLib import brepbndlib + + shape = scale_shape_to_unit_box(load_step_shape(one_step_fixture)) + box = Bnd_Box() + brepbndlib.Add(shape, box) + xmin, ymin, zmin, xmax, ymax, zmax = box.Get() + half = max(xmax - xmin, ymax - ymin, zmax - zmin) / 2.0 + # Largest half-extent normalised to ~1. + assert half == pytest.approx(1.0, abs=1e-3) diff --git a/ll_brepnet/tests/test_json_frontend.py b/ll_brepnet/tests/test_json_frontend.py new file mode 100644 index 0000000..bccde0a --- /dev/null +++ b/ll_brepnet/tests/test_json_frontend.py @@ -0,0 +1,54 @@ +"""Tests for the JSON topology front-end (no pythonocc required).""" + +from __future__ import annotations + +import numpy as np +import pytest + +from ll_brepnet.pipelines.extract_brepnet_data_from_json import BRepJsonExtractor + + +def _valid_topology(): + # Two faces, one shared edge -> two coedges (mates of each other). + return { + "coedge_to_next": [0, 1], + "coedge_to_prev": [0, 1], + "coedge_to_mate": [1, 0], + "coedge_to_face": [0, 1], + "coedge_to_edge": [0, 0], + "face_features": [[1.0] * 8, [0.0] * 8], + "edge_features": [[0.5] * 7], + } + + +def test_json_extractor_builds_arrays(): + arrays = BRepJsonExtractor(_valid_topology()).extract_arrays() + assert int(arrays["num_faces"]) == 2 + assert int(arrays["num_edges"]) == 1 + assert int(arrays["num_coedges"]) == 2 + # Grids default to zeros with the right shape when absent. + assert arrays["face_point_grids"].shape == (2, 7, 10, 10) + assert arrays["edge_point_grids"].shape == (1, 6, 10) + assert np.array_equal(arrays["coedge_to_mate"], np.array([1, 0])) + + +def test_json_extractor_missing_key_raises(): + topo = _valid_topology() + del topo["face_features"] + with pytest.raises(ValueError, match="face_features"): + BRepJsonExtractor(topo).extract_arrays() + + +def test_json_extractor_out_of_range_index_raises(): + topo = _valid_topology() + topo["coedge_to_face"] = [0, 5] # only 2 faces exist + with pytest.raises(ValueError, match="coedge_to_face"): + BRepJsonExtractor(topo).extract_arrays() + + +def test_json_extractor_process_writes_npz(tmp_path): + out = tmp_path / "solid.npz" + BRepJsonExtractor(_valid_topology()).process(out) + assert out.exists() + with np.load(out) as data: + assert int(data["num_coedges"]) == 2 diff --git a/ll_brepnet/tests/test_model.py b/ll_brepnet/tests/test_model.py new file mode 100644 index 0000000..d8ca679 --- /dev/null +++ b/ll_brepnet/tests/test_model.py @@ -0,0 +1,93 @@ +"""Tests for the LLBRepNet model: forward shape, finiteness, learning.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("torch") +pytest.importorskip("OCC") + +import torch # noqa: E402 +import torch.nn.functional as F # noqa: E402 + +from ll_brepnet.dataloaders.brep_dataset import BRepDataset, brep_collate_fn # noqa: E402 +from ll_brepnet.models.ll_brepnet import LLBRepNet # noqa: E402 +from ll_brepnet.models.uvnet_encoders import UVNetCurveEncoder, UVNetSurfaceEncoder # noqa: E402 + + +def _batch(prepared_dataset): + npz_dir, manifest = prepared_dataset + ds = BRepDataset(manifest, npz_dir, "training_set", label_dir=npz_dir) + if len(ds) == 0: + pytest.skip("empty training split") + return brep_collate_fn([ds[i] for i in range(len(ds))]) + + +def test_uvnet_encoder_shapes(): + surf = UVNetSurfaceEncoder(in_channels=7, out_dim=16) + curve = UVNetCurveEncoder(in_channels=6, out_dim=16) + assert surf(torch.randn(5, 7, 10, 10)).shape == (5, 16) + assert curve(torch.randn(5, 6, 10)).shape == (5, 16) + # Empty inputs are handled. + assert surf(torch.zeros(0, 7, 10, 10)).shape == (0, 16) + + +@pytest.mark.requires_pythonocc +def test_forward_shape_and_finite(prepared_dataset): + batch = _batch(prepared_dataset) + model = LLBRepNet( + num_classes=7, + num_layers=2, + hidden_dim=32, + entity_hidden=16, + surf_emb_dim=16, + curve_emb_dim=16, + ) + model.eval() + logits = model(batch) + assert logits.shape == (batch.num_faces, 7) + assert torch.isfinite(logits).all() + + +@pytest.mark.requires_pythonocc +def test_one_step_decreases_loss(prepared_dataset): + torch.manual_seed(0) + batch = _batch(prepared_dataset) + model = LLBRepNet( + num_classes=7, + num_layers=2, + hidden_dim=32, + entity_hidden=16, + surf_emb_dim=16, + curve_emb_dim=16, + ) + model.train() + opt = torch.optim.Adam(model.parameters(), lr=1e-2) + first = last = None + for _step in range(12): + opt.zero_grad() + loss = F.cross_entropy(model(batch), batch.labels, ignore_index=-1) + loss.backward() + opt.step() + if first is None: + first = float(loss) + last = float(loss) + assert last < first + + +@pytest.mark.requires_pythonocc +def test_predict_logits_per_solid(prepared_dataset): + batch = _batch(prepared_dataset) + model = LLBRepNet( + num_classes=7, + num_layers=2, + hidden_dim=32, + entity_hidden=16, + surf_emb_dim=16, + curve_emb_dim=16, + ) + out = model.predict_logits(batch) + assert len(out) == batch.num_solids + for (_stem, probs), (start, end) in zip(out, batch.split_batch): + assert probs.shape == (end - start, 7) + assert torch.allclose(probs.sum(dim=1), torch.ones(end - start), atol=1e-4) diff --git a/site/astro.config.mjs b/site/astro.config.mjs index 9947ac6..d7104cb 100644 --- a/site/astro.config.mjs +++ b/site/astro.config.mjs @@ -52,6 +52,7 @@ export default defineConfig({ { label: 'll_ocadr', items: [{ autogenerate: { directory: 'll_ocadr' } }] }, { label: 'll_gen', items: [{ autogenerate: { directory: 'll_gen' } }] }, { label: 'll_clouds', items: [{ autogenerate: { directory: 'll_clouds' } }] }, + { label: 'll_brepnet', items: [{ autogenerate: { directory: 'll_brepnet' } }] }, ], }, { label: 'Roadmap', items: [{ autogenerate: { directory: 'roadmap' } }] }, diff --git a/site/src/content/docs/ll_brepnet/installation.md b/site/src/content/docs/ll_brepnet/installation.md new file mode 100644 index 0000000..ba629b9 --- /dev/null +++ b/site/src/content/docs/ll_brepnet/installation.md @@ -0,0 +1,55 @@ +--- +title: ll_brepnet — Installation +description: Install ll_brepnet — PyTorch + pythonocc-core + occwl via conda, plus pytorch-lightning. +sidebar: + label: Installation + order: 2 +--- + +`ll_brepnet` needs PyTorch, `pythonocc-core`, and `occwl`, which on macOS must +come from conda to avoid the OpenMP/`libomp` conflict (see the repo `CLAUDE.md`). +It also uses `pytorch-lightning` for training and `cadling` for the shared B-Rep +extraction machinery. + +## Option 1 — reuse the `cadling` environment (recommended) + +The `cadling` conda environment already provides PyTorch, `pythonocc-core`, and +`occwl`. Add the training deps and install the package editable: + +```bash +conda activate cadling +pip install pytorch-lightning tensorboard +pip install -e ./ll_brepnet +``` + +## Option 2 — standalone environment + +```bash +conda env create -f ll_brepnet/environment.yaml +conda activate ll-brepnet +``` + +## Verify + +```bash +python -c "import ll_brepnet; print(ll_brepnet.__version__)" +# 0.1.0 +``` + +Run the test suite (skips automatically without pythonocc / torch): + +```bash +pytest ll_brepnet/tests -q # fast tests +pytest ll_brepnet/tests -q -m "" # include the slow end-to-end test +``` + +## Dependencies + +| Dependency | Source | Why | +|---|---|---| +| `pytorch` | conda-forge | model + training (conda only, OpenMP safety) | +| `pythonocc-core` | conda-forge | STEP loading + B-Rep traversal | +| `occwl` | pip | UV-grid sampling (`uvgrid` / `ugrid`) | +| `pytorch-lightning`, `torchmetrics`, `tensorboard` | pip | training loop, mIoU/accuracy, logging | +| `cadling` | editable (monorepo) | coedge extraction + reused encoder | +| `scikit-learn` | conda/pip | train/val/test splitting | diff --git a/site/src/content/docs/ll_brepnet/overview.md b/site/src/content/docs/ll_brepnet/overview.md new file mode 100644 index 0000000..f523f00 --- /dev/null +++ b/site/src/content/docs/ll_brepnet/overview.md @@ -0,0 +1,87 @@ +--- +title: ll_brepnet — Overview +description: A B-Rep face-graph neural network for CAD solid-model segmentation — coedge message passing + UV-grid geometry, trained on the Fusion 360 Gallery set. +sidebar: + label: Overview + order: 1 + badge: + text: Trained + variant: success +--- + +**ll_brepnet** is a B-Rep **face-segmentation** neural network. It operates +directly on the boundary representation of a CAD solid — faces and edges +connected through oriented *coedges* (half-edges carrying next / previous / mate +/ parent-face / parent-edge adjacency) — and fuses that topology with UV-grid +surface and curve geometry to predict a semantic segment label for every face. + +It is an **independent, MIT-licensed** package built on the toolkit's own B-Rep +machinery (`cadling`). It is *inspired by* BRepNet +([arXiv:2104.00706](https://arxiv.org/abs/2104.00706)) and UV-Net, but contains +no code from those projects — see `ll_brepnet/ATTRIBUTION.md`. + +## Modules + +| Module | Responsibility | +|---|---| +| `ll_brepnet.pipelines.extract_brepnet_data_from_step` | STEP → unit-box scale → coedge graph + per-face/edge features + UV-grids → `.npz` | +| `ll_brepnet.pipelines.extract_brepnet_data_from_json` | Same record from a precomputed JSON topology | +| `ll_brepnet.pipelines.build_dataset_file` / `quickstart` | Manifest building (split + train-only standardization); Fusion 360 orchestration | +| `ll_brepnet.dataloaders` | `BRepDataset`, offset-aware `brep_collate_fn`, `MaxNumFacesSampler`, `BRepDataModule` | +| `ll_brepnet.models` | UV-Net surface/curve encoders + the `LLBRepNet` LightningModule | +| `ll_brepnet.train` / `ll_brepnet.eval` | pytorch-lightning training; folder/checkpoint inference → per-face logits | + +## Architecture + +```text +STEP solid + → unit-box scale → coedge graph (next/prev/mate/face/edge incidence) + → per-face features (surface-type one-hot + area) + face UV-grid [7,U,V] + → per-edge features (curve-type one-hot + length + convexity) + edge U-grid [6,U] + │ + ▼ + UV-Net surface/curve encoders ⊕ scalar features + │ (gather per coedge: face_repr[c2f] ⊕ edge_repr[c2e] ⊕ reversed) + ▼ + coedge message passing (BRepNetEncoder, reused from cadling) → coedge→face mean pool + ▼ + per-face segmentation head → [num_faces, num_classes] +``` + +## Results + +Trained on a real subset of the **Fusion 360 Gallery segmentation dataset +(s2.0.0)** — 3,400 train / 600 validation / 800 test solids, the official +train/test partition, the official 8 manufacturing-feature classes, 35 epochs +on CPU. + +**Held-out test split (800 real solids):** + +| Metric | Value | +|---|---| +| **mean IoU (macro)** | **0.709** | +| accuracy | 0.912 | + +Per-class IoU: + +| Class | IoU | | Class | IoU | +|---|---|---|---|---| +| Fillet | 0.94 | | CutEnd | 0.71 | +| ExtrudeSide | 0.89 | | RevolveSide | 0.66 | +| ExtrudeEnd | 0.86 | | CutSide | 0.66 | +| Chamfer | 0.84 | | RevolveEnd | 0.11 | + +These are **real, reproducible numbers** (see **Usage**), competitive with the +BRepNet paper's reported ~0.65–0.72 mIoU on s2.0.0 — achieved here with the +MIT-clean reused-coedge-encoder architecture rather than the paper's kernel +convolution. + +:::caution[Honest scope] +This checkpoint trains on a **subset** of the full 35,680-solid dataset and for +a bounded number of CPU epochs. The rare **RevolveEnd** class (IoU 0.11) is +under-represented in this subset and barely learned — training on the full set +(and/or more epochs / a GPU) would raise the tail classes. The package ships the +training/eval code, not a production checkpoint. +::: + +Use the sidebar for **Installation** and **Usage**. diff --git a/site/src/content/docs/ll_brepnet/usage.md b/site/src/content/docs/ll_brepnet/usage.md new file mode 100644 index 0000000..557fc7b --- /dev/null +++ b/site/src/content/docs/ll_brepnet/usage.md @@ -0,0 +1,97 @@ +--- +title: ll_brepnet — Usage +description: Prepare the Fusion 360 dataset, train the segmentation model, and run inference on STEP files. +sidebar: + label: Usage + order: 3 +--- + +The pipeline is **STEP → `.npz` records → train → segment**. Everything below +uses the real Fusion 360 Gallery segmentation dataset (s2.0.0). + +## 1. Get the dataset + +```bash +curl https://fusion-360-gallery-dataset.s3-us-west-2.amazonaws.com/segmentation/s2.0.0/s2.0.0.zip -o s2.0.0.zip +unzip s2.0.0.zip # → s2.0.0/{breps/step, breps/seg, train_test.json, segment_names.json} +``` + +## 2. Prepare records + manifest + +`quickstart` extracts each STEP file to a `.npz`, copies the matching `.seg` +labels, and writes a `dataset.json` (official train/test split, train-only +feature standardization, the 8 class names). Use `--train-subset` / +`--test-subset` to work on a subset, or omit them for the full 35,680 solids. + +```bash +python -m ll_brepnet.pipelines.quickstart \ + --dataset-dir ./s2.0.0 \ + --output-dir ./processed \ + --num-workers 12 \ + --train-subset 4000 --test-subset 800 +``` + +## 3. Train + +```bash +python -m ll_brepnet.train \ + --dataset-file ./processed/dataset.json \ + --dataset-dir ./processed \ + --label-dir ./processed \ + --max-epochs 35 --num-layers 4 --hidden-dim 128 \ + --learning-rate 0.002 --log-dir ./logs +``` + +Checkpoints (best by validation mIoU) and TensorBoard logs land under +`--log-dir`. **This exact command produced test mIoU ≈ 0.709 / accuracy ≈ 0.912** +on the 800-solid test split (see [Overview](/ll_toolkit/ll_brepnet/overview/)). + +## 4. Segment STEP files + +```bash +python -m ll_brepnet.eval.evaluate \ + --step-dir ./s2.0.0/breps/step \ + --model ./logs/ll_brepnet/version_0/checkpoints/best.ckpt \ + --dataset-file ./processed/dataset.json \ + --output-dir ./predictions +``` + +Each input solid gets a `.logits` file: one row per face, one column per +class. The predicted segment of face *i* is the `argmax` of row *i*. + +## Python API + +```python +from ll_brepnet import ( + extract_brepnet_data_from_step, # STEP folder → .npz records + prepare_fusion360, # Fusion 360 dir → records + dataset.json + BRepDataModule, # LightningDataModule + LLBRepNet, # the model (a LightningModule) + evaluate_folder, # checkpoint + STEP folder → per-face logits +) + +# Inference from an existing checkpoint: +written = evaluate_folder( + step_dir="./s2.0.0/breps/step", + model_ckpt="./logs/.../best.ckpt", + dataset_file="./processed/dataset.json", + output_dir="./predictions", + num_workers=8, +) +``` + +The model also exposes `predict_logits(batch)` → `[(file_stem, per-face softmax)]` +for programmatic, per-solid inference. + +## Alternate input: precomputed JSON topology + +If you already have the coedge topology + features from another tool, skip the +OpenCascade extraction and feed JSON directly: + +```bash +python -m ll_brepnet.pipelines.extract_brepnet_data_from_json \ + --json my_solid.json --output my_solid.npz +``` + +See `BRepJsonExtractor` for the schema (coedge incidence arrays + per-face/edge +feature matrices; UV-grids optional). diff --git a/site/src/content/docs/roadmap/ll_brepnet.md b/site/src/content/docs/roadmap/ll_brepnet.md index 0b293bc..7e91aa5 100644 --- a/site/src/content/docs/roadmap/ll_brepnet.md +++ b/site/src/content/docs/roadmap/ll_brepnet.md @@ -1,55 +1,32 @@ --- -title: ll_brepnet (Planned) -description: A planned B-Rep face-graph neural network. Not yet implemented — the package is currently an empty scaffold. +title: ll_brepnet (Implemented) +description: The planned B-Rep face-graph neural network is now implemented and trained — see the ll_brepnet package docs. sidebar: - label: ll_brepnet + label: ll_brepnet → done order: 1 badge: - text: Planned - variant: caution + text: Done + variant: success --- -:::caution[Not implemented yet] -`ll_brepnet` is **planned, not built**. The directory exists in the repository, -but every file in it is currently empty (0 bytes) — including `pyproject.toml`, -`requirements.txt`, and all of its Python modules. It is **not** installable, -importable, or runnable today, and it is intentionally **not** listed among the -toolkit's shipping packages. +:::tip[Implemented] +`ll_brepnet` is **no longer a roadmap item — it is built, trained, and +documented.** The package extracts the coedge graph + UV-grid geometry from STEP +solids and trains a face-segmentation model end-to-end. -This page exists so the project's status is honest: the toolkit documents what -its code actually does. When `ll_brepnet` has a real implementation, it will get -full Overview / Installation / Usage / API Reference pages like the other -packages — and a roadmap entry will no longer be needed. +On a real subset of the **Fusion 360 Gallery segmentation dataset** (official +8 manufacturing-feature classes, 800-solid held-out test split) it reaches +**test mIoU ≈ 0.709 / accuracy ≈ 0.912** — competitive with the BRepNet paper's +reported ~0.65–0.72 mIoU, achieved with an MIT-clean architecture. ::: -## What it is intended to be +See the package documentation: -Based on the scaffolded module layout, `ll_brepnet` is intended to be a **B-Rep -face-graph neural network** in the lineage of UV-Net / BRepNet — models that -operate directly on the boundary-representation graph of a CAD solid (faces as -nodes, edges as connections) with UV-grid surface features. +- [Overview](/ll_toolkit/ll_brepnet/overview/) — architecture, results, per-class IoU +- [Installation](/ll_toolkit/ll_brepnet/installation/) +- [Usage](/ll_toolkit/ll_brepnet/usage/) — prepare data, train, and segment STEP files -The planned structure (currently empty placeholders) sketches: - -| Area | Planned modules | -|---|---| -| Data loading | `dataloaders/brep_dataset.py`, `dataloaders/max_num_faces_loader.py` | -| Models | `models/ll_brepnet.py`, `models/uvnet_encoders.py` | -| Pipelines | `pipelines/extract_brepnet_data_from_step.py`, `extract_brepnet_data_from_json.py`, `entity_mapper.py`, `build_dataset_file.py` | -| Evaluation | `eval/evaluate.py` | - -## How it would fit the toolkit - -A B-Rep face-graph model complements the existing neural packages: - -- [cadling](/ll_toolkit/cadling/overview/) already builds a B-Rep topology graph - during parsing — a natural input source. -- [geotoken](/ll_toolkit/geotoken/overview/) already provides B-Rep - graph tokenization (`GraphTokenizer`). -- [ll_stepnet](/ll_toolkit/ll_stepnet/overview/) covers STEP-text + topology - fusion; `ll_brepnet` would specialize in pure face-graph learning. - -## Status - -There is no code, no tests, and no checkpoints. Follow the repository for -updates: . +It is an independent, MIT-licensed implementation built on the toolkit's own +B-Rep machinery (`cadling`); inspired by BRepNet +([arXiv:2104.00706](https://arxiv.org/abs/2104.00706)) and UV-Net but containing +no code from those projects.