Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Writeup.pdf filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.pyc
+*.pyo
+.DS_Store
+*.egg-info/
+.ipynb_checkpoints/

LICENSE.md

@@ -0,0 +1,408 @@
+Attribution-NonCommercial 4.0 International
+
+=======================================================================
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright
+and certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+     Considerations for licensors: Our public licenses are
+     intended for use by those authorized to give the public
+     permission to use material in ways otherwise restricted by
+     copyright and certain other rights. Our licenses are
+     irrevocable. Licensors should read and understand the terms
+     and conditions of the license they choose before applying it.
+     Licensors should also secure all rights necessary before
+     applying our licenses so that the public can reuse the
+     material as expected. Licensors should clearly mark any
+     material not subject to the license. This includes other CC-
+     licensed material, or material used under an exception or
+     limitation to copyright. More considerations for licensors:
+    wiki.creativecommons.org/Considerations_for_licensors
+
+     Considerations for the public: By using one of our public
+     licenses, a licensor grants the public permission to use the
+     licensed material under specified terms and conditions. If
+     the licensor's permission is not necessary for any reason--for
+     example, because of any applicable exception or limitation to
+     copyright--then that use is not regulated by the license. Our
+     licenses grant only permissions under copyright and certain
+     other rights that a licensor has authority to grant. Use of
+     the licensed material may still be restricted for other
+     reasons, including because others have copyright or other
+     rights in the material. A licensor may make special requests,
+     such as asking that all changes be marked or described.
+     Although not required by our licenses, you are encouraged to
+     respect those requests where reasonable. More considerations
+     for the public:
+    wiki.creativecommons.org/Considerations_for_licensees
+
+=======================================================================
+
+Creative Commons Attribution-NonCommercial 4.0 International Public
+License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial 4.0 International Public License ("Public
+License"). To the extent this Public License may be interpreted as a
+contract, You are granted the Licensed Rights in consideration of Your
+acceptance of these terms and conditions, and the Licensor grants You
+such rights in consideration of benefits the Licensor receives from
+making the Licensed Material available under these terms and
+conditions.
+
+
+Section 1 -- Definitions.
+
+  a. Adapted Material means material subject to Copyright and Similar
+     Rights that is derived from or based upon the Licensed Material
+     and in which the Licensed Material is translated, altered,
+     arranged, transformed, or otherwise modified in a manner requiring
+     permission under the Copyright and Similar Rights held by the
+     Licensor. For purposes of this Public License, where the Licensed
+     Material is a musical work, performance, or sound recording,
+     Adapted Material is always produced where the Licensed Material is
+     synched in timed relation with a moving image.
+
+  b. Adapter's License means the license You apply to Your Copyright
+     and Similar Rights in Your contributions to Adapted Material in
+     accordance with the terms and conditions of this Public License.
+
+  c. Copyright and Similar Rights means copyright and/or similar rights
+     closely related to copyright including, without limitation,
+     performance, broadcast, sound recording, and Sui Generis Database
+     Rights, without regard to how the rights are labeled or
+     categorized. For purposes of this Public License, the rights
+     specified in Section 2(b)(1)-(2) are not Copyright and Similar
+     Rights.
+  d. Effective Technological Measures means those measures that, in the
+     absence of proper authority, may not be circumvented under laws
+     fulfilling obligations under Article 11 of the WIPO Copyright
+     Treaty adopted on December 20, 1996, and/or similar international
+     agreements.
+
+  e. Exceptions and Limitations means fair use, fair dealing, and/or
+     any other exception or limitation to Copyright and Similar Rights
+     that applies to Your use of the Licensed Material.
+
+  f. Licensed Material means the artistic or literary work, database,
+     or other material to which the Licensor applied this Public
+     License.
+
+  g. Licensed Rights means the rights granted to You subject to the
+     terms and conditions of this Public License, which are limited to
+     all Copyright and Similar Rights that apply to Your use of the
+     Licensed Material and that the Licensor has authority to license.
+
+  h. Licensor means the individual(s) or entity(ies) granting rights
+     under this Public License.
+
+  i. NonCommercial means not primarily intended for or directed towards
+     commercial advantage or monetary compensation. For purposes of
+     this Public License, the exchange of the Licensed Material for
+     other material subject to Copyright and Similar Rights by digital
+     file-sharing or similar means is NonCommercial provided there is
+     no payment of monetary compensation in connection with the
+     exchange.
+
+  j. Share means to provide material to the public by any means or
+     process that requires permission under the Licensed Rights, such
+     as reproduction, public display, public performance, distribution,
+     dissemination, communication, or importation, and to make material
+     available to the public including in ways that members of the
+     public may access the material from a place and at a time
+     individually chosen by them.
+
+  k. Sui Generis Database Rights means rights other than copyright
+     resulting from Directive 96/9/EC of the European Parliament and of
+     the Council of 11 March 1996 on the legal protection of databases,
+     as amended and/or succeeded, as well as other essentially
+     equivalent rights anywhere in the world.
+
+  l. You means the individual or entity exercising the Licensed Rights
+     under this Public License. Your has a corresponding meaning.
+
+
+Section 2 -- Scope.
+
+  a. License grant.
+
+       1. Subject to the terms and conditions of this Public License,
+          the Licensor hereby grants You a worldwide, royalty-free,
+          non-sublicensable, non-exclusive, irrevocable license to
+          exercise the Licensed Rights in the Licensed Material to:
+
+            a. reproduce and Share the Licensed Material, in whole or
+               in part, for NonCommercial purposes only; and
+
+            b. produce, reproduce, and Share Adapted Material for
+               NonCommercial purposes only.
+
+       2. Exceptions and Limitations. For the avoidance of doubt, where
+          Exceptions and Limitations apply to Your use, this Public
+          License does not apply, and You do not need to comply with
+          its terms and conditions.
+
+       3. Term. The term of this Public License is specified in Section
+          6(a).
+
+       4. Media and formats; technical modifications allowed. The
+          Licensor authorizes You to exercise the Licensed Rights in
+          all media and formats whether now known or hereafter created,
+          and to make technical modifications necessary to do so. The
+          Licensor waives and/or agrees not to assert any right or
+          authority to forbid You from making technical modifications
+          necessary to exercise the Licensed Rights, including
+          technical modifications necessary to circumvent Effective
+          Technological Measures. For purposes of this Public License,
+          simply making modifications authorized by this Section 2(a)
+          (4) never produces Adapted Material.
+
+       5. Downstream recipients.
+
+            a. Offer from the Licensor -- Licensed Material. Every
+               recipient of the Licensed Material automatically
+               receives an offer from the Licensor to exercise the
+               Licensed Rights under the terms and conditions of this
+               Public License.
+
+            b. No downstream restrictions. You may not offer or impose
+               any additional or different terms or conditions on, or
+               apply any Effective Technological Measures to, the
+               Licensed Material if doing so restricts exercise of the
+               Licensed Rights by any recipient of the Licensed
+               Material.
+
+       6. No endorsement. Nothing in this Public License constitutes or
+          may be construed as permission to assert or imply that You
+          are, or that Your use of the Licensed Material is, connected
+          with, or sponsored, endorsed, or granted official status by,
+          the Licensor or others designated to receive attribution as
+          provided in Section 3(a)(1)(A)(i).
+
+  b. Other rights.
+
+       1. Moral rights, such as the right of integrity, are not
+          licensed under this Public License, nor are publicity,
+          privacy, and/or other similar personality rights; however, to
+          the extent possible, the Licensor waives and/or agrees not to
+          assert any such rights held by the Licensor to the limited
+          extent necessary to allow You to exercise the Licensed
+          Rights, but not otherwise.
+
+       2. Patent and trademark rights are not licensed under this
+          Public License.
+
+       3. To the extent possible, the Licensor waives any right to
+          collect royalties from You for the exercise of the Licensed
+          Rights, whether directly or through a collecting society
+          under any voluntary or waivable statutory or compulsory
+          licensing scheme. In all other cases the Licensor expressly
+          reserves any right to collect such royalties, including when
+          the Licensed Material is used other than for NonCommercial
+          purposes.
+
+
+Section 3 -- License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+  a. Attribution.
+
+       1. If You Share the Licensed Material (including in modified
+          form), You must:
+
+            a. retain the following if it is supplied by the Licensor
+               with the Licensed Material:
+
+                 i. identification of the creator(s) of the Licensed
+                    Material and any others designated to receive
+                    attribution, in any reasonable manner requested by
+                    the Licensor (including by pseudonym if
+                    designated);
+
+                ii. a copyright notice;
+
+               iii. a notice that refers to this Public License;
+
+                iv. a notice that refers to the disclaimer of
+                    warranties;
+
+                 v. a URI or hyperlink to the Licensed Material to the
+                    extent reasonably practicable;
+
+            b. indicate if You modified the Licensed Material and
+               retain an indication of any previous modifications; and
+
+            c. indicate the Licensed Material is licensed under this
+               Public License, and include the text of, or the URI or
+               hyperlink to, this Public License.
+
+       2. You may satisfy the conditions in Section 3(a)(1) in any
+          reasonable manner based on the medium, means, and context in
+          which You Share the Licensed Material. For example, it may be
+          reasonable to satisfy the conditions by providing a URI or
+          hyperlink to a resource that includes the required
+          information.
+
+       3. If requested by the Licensor, You must remove any of the
+          information required by Section 3(a)(1)(A) to the extent
+          reasonably practicable.
+
+       4. If You Share Adapted Material You produce, the Adapter's
+          License You apply must not prevent recipients of the Adapted
+          Material from complying with this Public License.
+
+
+Section 4 -- Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that
+apply to Your use of the Licensed Material:
+
+  a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+     to extract, reuse, reproduce, and Share all or a substantial
+     portion of the contents of the database for NonCommercial purposes
+     only;
+
+  b. if You include all or a substantial portion of the database
+     contents in a database in which You have Sui Generis Database
+     Rights, then the database in which You have Sui Generis Database
+     Rights (but not its individual contents) is Adapted Material; and
+
+  c. You must comply with the conditions in Section 3(a) if You Share
+     all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not
+replace Your obligations under this Public License where the Licensed
+Rights include other Copyright and Similar Rights.
+
+
+Section 5 -- Disclaimer of Warranties and Limitation of Liability.
+
+  a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
+     EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
+     AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
+     ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
+     IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
+     WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
+     PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
+     ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
+     KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
+     ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
+
+  b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
+     TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
+     NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
+     INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
+     COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
+     USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
+     ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
+     DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
+     IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
+
+  c. The disclaimer of warranties and limitation of liability provided
+     above shall be interpreted in a manner that, to the extent
+     possible, most closely approximates an absolute disclaimer and
+     waiver of all liability.
+
+
+Section 6 -- Term and Termination.
+
+  a. This Public License applies for the term of the Copyright and
+     Similar Rights licensed here. However, if You fail to comply with
+     this Public License, then Your rights under this Public License
+     terminate automatically.
+
+  b. Where Your right to use the Licensed Material has terminated under
+     Section 6(a), it reinstates:
+
+       1. automatically as of the date the violation is cured, provided
+          it is cured within 30 days of Your discovery of the
+          violation; or
+
+       2. upon express reinstatement by the Licensor.
+
+     For the avoidance of doubt, this Section 6(b) does not affect any
+     right the Licensor may have to seek remedies for Your violations
+     of this Public License.
+
+  c. For the avoidance of doubt, the Licensor may also offer the
+     Licensed Material under separate terms or conditions or stop
+     distributing the Licensed Material at any time; however, doing so
+     will not terminate this Public License.
+
+  d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+     License.
+
+
+Section 7 -- Other Terms and Conditions.
+
+  a. The Licensor shall not be bound by any additional or different
+     terms or conditions communicated by You unless expressly agreed.
+
+  b. Any arrangements, understandings, or agreements regarding the
+     Licensed Material not stated herein are separate from and
+     independent of the terms and conditions of this Public License.
+
+
+Section 8 -- Interpretation.
+
+  a. For the avoidance of doubt, this Public License does not, and
+     shall not be interpreted to, reduce, limit, restrict, or impose
+     conditions on any use of the Licensed Material that could lawfully
+     be made without permission under this Public License.
+
+  b. To the extent possible, if any provision of this Public License is
+     deemed unenforceable, it shall be automatically reformed to the
+     minimum extent necessary to make it enforceable. If the provision
+     cannot be reformed, it shall be severed from this Public License
+     without affecting the enforceability of the remaining terms and
+     conditions.
+
+  c. No term or condition of this Public License will be waived and no
+     failure to comply consented to unless expressly agreed to by the
+     Licensor.
+
+  d. Nothing in this Public License constitutes or may be interpreted
+     as a limitation upon, or waiver of, any privileges and immunities
+     that apply to the Licensor or You, including from the legal
+     processes of any jurisdiction or authority.
+
+=======================================================================
+
+Creative Commons is not a party to its public
+licenses. Notwithstanding, Creative Commons may elect to apply one of
+its public licenses to material it publishes and in those instances
+will be considered the “Licensor.” The text of the Creative Commons
+public licenses is dedicated to the public domain under the CC0 Public
+Domain Dedication. Except for the limited purpose of indicating that
+material is shared under a Creative Commons public license or as
+otherwise permitted by the Creative Commons policies published at
+creativecommons.org/policies, Creative Commons does not authorize the
+use of the trademark "Creative Commons" or any other trademark or logo
+of Creative Commons without its prior written consent including,
+without limitation, in connection with any unauthorized modifications
+to any of its public licenses or any other arrangements,
+understandings, or agreements concerning use of licensed material. For
+the avoidance of doubt, this paragraph does not form part of the
+public licenses.
+
+Creative Commons may be contacted at creativecommons.org.
+

README.md

@@ -0,0 +1,66 @@
+# S23DR 2026 — Iterating on the Baseline (5th place)
+
+Solution for the [S23DR 2026](https://huggingface.co/spaces/usm3d/S23DR2026)
+structured 3D wireframe reconstruction challenge. Private leaderboard:
+**0.5388 HSS, 5th place**, with a single 8.85M-parameter Perceiver.
+
+The submitted entry is the **raw 8k Perceiver**: a learned segment model over a
+fused, priority-sampled COLMAP/depth point cloud, with no hand-crafted
+post-processing. The repository also contains the hand-crafted pipeline and the
+classifier-gated hybrid that we used earlier in the competition, and the
+training/repro scripts behind them. See the accompanying write-up for the full
+account.
+
+## Run inference
+
+```bash
+pip install -r requirements.txt
+python script.py
+```
+
+The challenge harness provides `params.json`, downloads the dataset, runs
+`script.py`, and reads the resulting `submission.json`
+(`{order_id, wf_vertices, wf_edges}` per scene). `script.py` loads
+`checkpoint_8192.pt` and runs the raw 8k model (`CONF_THRESH=0.5`, no seam,
+no augments).
+
+## Layout
+
+```
+script.py                 raw 8k inference (the submitted entry)
+Writeup.pdf               method write-up (full account of the solution)
+checkpoint_8192.pt        8k Perceiver weights (the 5th-place model)
+checkpoint.pt             organizers' 4k Perceiver (curriculum start point)
+solution.py               hand-crafted geometric pipeline
+edge_classifier.py        PointNet edge classifier (hybrid augment)
+vertex_refiner.py         PointNet vertex classifier (hybrid augment)
+pnet_class_2026.pth       edge classifier weights
+vertex_refiner.pth        vertex classifier weights
+s23dr_2026_example/       model + training package (Perceiver, tokenizer, losses, train.py)
+configs/                  training config (base.json)
+REPRODUCE.md              recipe for the resolution curriculum (2k -> 4k -> 8k)
+training/                 data-generation and training scripts (see below)
+experiments/ptv3/         the Point Transformer V3 encoder experiment
+                          (negative result; trained model + logs + code)
+```
+
+## Reproducing
+
+**The 8k model.** Follow `REPRODUCE.md`: train the Perceiver from scratch at
+2048 points, then fine-tune at 4096 and 8192 points on the organizers' released
+sampled point clouds. Our contribution is the 4k→8k stage; the released 4k
+checkpoint (`checkpoint.pt`) is the starting point.
+
+**Classifier augments.** `training/gen_edge_dataset.py` and
+`training/gen_vertex_dataset.py` build the per-candidate patch datasets from the
+hand-crafted pipeline's predictions; `training/train_edge_classifier_2026.py`
+and `training/train_vertex_refiner_2026.py` train the PointNet classifiers.
+
+**Router (negative result).** `training/train_routing_gbt.py` with
+`training/oracle_sources_{train,validation}.json` reproduces the gradient-boosted
+per-scene router; it recovered only 4.5% of the per-scene oracle ceiling.
+
+Some scripts under `training/` contain absolute paths from the original training
+environment and expect the repository root on `PYTHONPATH`
+(e.g. `PYTHONPATH=. python training/train_edge_classifier_2026.py`); adapt the
+paths to your setup.

REPRODUCE.md

@@ -0,0 +1,121 @@
+# Reproducing the 8k Perceiver (private HSS 0.5388, 5th place)
+
+The submitted entry is `checkpoint_8192.pt`, the raw 8k Perceiver. It is the end
+of a resolution curriculum: the organizers' 2048 → 4096 baseline produces
+`checkpoint.pt`, which we then fine-tune at 8192 points to produce
+`checkpoint_8192.pt`. The architecture is identical at every stage — only the
+input point budget grows.
+
+The Perceiver training code is the organizers' package, bundled here under
+`s23dr_2026_example/` (`train.py`, `tokenizer.py`, `model.py`, `losses.py`).
+
+## Inference
+
+```bash
+pip install -r requirements.txt
+python script.py
+```
+
+`script.py` loads `checkpoint_8192.pt`, fuses and priority-samples each scene to
+8192 points (6144 COLMAP + 2048 depth), runs the Perceiver, and writes
+`submission.json` (`{order_id, wf_vertices, wf_edges}` per scene). No seam, no
+augments (`CONF_THRESH=0.5`).
+
+## Architecture
+
+Every stage shares one config (`configs/base.json`):
+
+```
+Perceiver: hidden=256, ff=1024, latent_tokens=256, latent_layers=7
+  encoder_layers=4, decoder_layers=3, cross_attn_interval=4
+  num_heads=4, kv_heads_cross=2, kv_heads_self=2
+  qk_norm=True (L2), rms_norm=True, dropout=0.1
+  segments=64, segment_param=midpoint_dir_len, segment_conf=True
+  behind_emb_dim=8, vote_features=True, activation=gelu
+```
+
+## Curriculum
+
+The Perceiver is trained by a staged resolution curriculum: from scratch at
+2048 points, then fine-tuned at 4096 and 8192. Each resolution step adds 45k
+fine-tuning steps, the last 20k a linear cooldown, at a gentle learning rate
+(3e-5) that preserves the lower-resolution representations. Stages A and B are
+the organizers' baseline; stage C is ours. The generic invocation is:
+
+```bash
+python -m s23dr_2026_example.train \
+    --cache-dir hf://usm3d/s23dr-2026-sampled_<N>:train \
+    --seq-len <N> [--resume <previous_checkpoint>] --aug-rotate --aug-flip
+```
+
+### Stage A — 2048, from scratch  *(organizers)*
+
+```
+Data:   sampled_2048_v2:train
+Steps:  0 -> 125,000     LR: 3e-4, warmup 10,000     Batch: 32
+Loss:   sinkhorn (eps=0.1, iters=20, dustbin=0.3) + conf (weight 0.1)
+Seed:   353
+```
+
+Trains the Perceiver from random init. The 2048 budget keeps the train/val gap
+low; training directly at high resolution overfits. **Public test HSS 0.4273.**
+
+### Stage B — 4096 fine-tune + cooldown  *(organizers)* → `checkpoint.pt`
+
+```
+Resume: Stage A          Data: sampled_4096_v2:train
+Steps:  125,000 -> 170,000  (45k; last 20k linear cooldown)
+LR:     3e-5             Batch: 64
+```
+
+Switches the input to 4096 points; the gentle LR adapts without disturbing the
+learned representation (LR > 1e-4 forgets it). The result is the organizers'
+released 4k checkpoint, `checkpoint.pt` — **public test HSS 0.4470**.
+
+### Stage C — 8192 fine-tune + cooldown  *(ours)* → `checkpoint_8192.pt`
+
+```
+Resume: checkpoint.pt    Data: organizers' released 8k samples
+Input:  8192 points = 6144 COLMAP + 2048 depth
+Steps:  170,000 -> 215,000  (45k; last 20k linear cooldown)
+LR:     3e-5             Batch: 64
+```
+
+Resumes the 4k checkpoint and continues the same gentle fine-tune at 8192 points.
+This is the only stage we contribute, and the single largest gain in the system:
+doubling the input from 4096 to 8192 points lifts the raw model from 0.4470 to
+**0.5004** public HSS (a +0.053 jump, larger than the entire hand-crafted hybrid
+contributes at 4k). Retuning the confidence threshold from 0.5 to 0.65 reaches
+our public best of 0.5009; the submitted entry uses 0.5. This is
+`checkpoint_8192.pt` — **public test HSS 0.5004, private 0.5388 (5th place)**.
+
+## Data
+
+The organizers publish pre-sampled point clouds (`sampled_2048` / `sampled_4096`
+/ `sampled_8192`) on the Hugging Face Hub; the curriculum above trains directly
+on them. `training/gen_sampled_16384.py` regenerates a sampled set at an
+arbitrary `--seq-len` from the organizers' `cached_full_pcd_v2` cache, in the
+same format — we used it to try a 16384 budget (see below).
+
+## Negative results (kept for the record, not on the inference path)
+
+- **16384 resolution** regressed relative to 8192. The generator is
+  `training/gen_sampled_16384.py`; the model was not shipped.
+- **Per-scene router** (`training/train_routing_gbt.py` with
+  `training/oracle_sources_{train,validation}.json`): a gradient-boosted router
+  choosing 4k / 8k / hand-crafted per scene recovered only ~4.5% of the
+  per-scene oracle ceiling.
+- **Point Transformer V3 encoder** (`experiments/ptv3/`): a stronger encoder,
+  too slow for the T4 2-hour budget. Trained checkpoint, logs, and code are
+  archived there.
+
+## Evaluation sets
+
+Two splits matter for the numbers above:
+
+- **Public test** — the competition harness scores `submission.json` against the
+  hidden test set and posts to the leaderboard. Every HSS in this document
+  (4k 0.4470, 8k 0.5004, private 0.5388) is a public/private test number.
+- **Dev val** — the tail of the published training set, used during development
+  to pick checkpoints. We do not quote dev-val numbers here, since the
+  leaderboard scores are the ones that decide the entry.

Writeup.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8854aca9ffbd1a6e6d1a97f95641f99b0cd352eeacc617b26abbe04583797065
+size 189408

checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1296423a1a2e603ba55860d8ef8fa3a861764a7bbc3de96b776fca59cf5b11ab
+size 106429791

checkpoint_8192.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f7283074ed44634770d1b5d7724cb08dac429df03d218cdfa0dcc72e75278bc
+size 106429599

configs/base.json

@@ -0,0 +1,39 @@
+{
+    "arch": "perceiver",
+    "segments": 64,
+    "hidden": 256,
+    "ff": 1024,
+    "num_heads": 4,
+    "kv_heads_cross": 2,
+    "kv_heads_self": 2,
+    "latent_tokens": 256,
+    "latent_layers": 7,
+    "decoder_layers": 3,
+    "cross_attn_interval": 4,
+    "encoder_layers": 4,
+    "behind_emb_dim": 8,
+    "dropout": 0.1,
+    "activation": "gelu",
+    "rms_norm": true,
+    "qk_norm": true,
+    "qk_norm_type": "l2",
+    "segment_param": "midpoint_dir_len",
+    "segment_conf": true,
+    "vote_features": true,
+
+    "adam_betas": "0.9,0.95",
+    "weight_decay": 0.01,
+    "warmup": 10000,
+    "varifold_weight": 0.0,
+    "sinkhorn_weight": 1.0,
+    "sinkhorn_eps": 0.1,
+    "sinkhorn_iters": 20,
+    "sinkhorn_dustbin": 0.3,
+    "conf_weight": 0.1,
+    "conf_mode": "sinkhorn",
+    "conf_head_wd": 0.1,
+
+    "aug_rotate": true,
+    "aug_flip": true,
+    "seed": 353
+}

edge_classifier.py

@@ -0,0 +1,208 @@
+"""Edge classifier and classifier-gated edge augmentation.
+
+Contains:
+  - ClassificationPointNet: PointNet binary classifier on 6D cylindrical patches
+  - colmap_points_xyz_rgb / build_edge_patch_6d: patch builder
+  - score_edges_batched: per-sample batched inference
+  - augment_hybrid_with_filtered_hc: import high-confidence handcrafted edges
+"""
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# Cylinder patch geometry (must match the training-time patch generation).
+CYL_RADIUS = 0.5
+CYL_EXT = 0.25            # extension at each end
+MAX_PATCH_POINTS = 1024
+AUGMENT_DEDUP_RADIUS = 0.3
+AUGMENT_THRESHOLD = 0.55
+
+
+class ClassificationPointNet(nn.Module):
+    """PointNet binary classifier on 6D point cloud patches (xyz + rgb)."""
+
+    def __init__(self, input_dim=6, max_points=1024):
+        super().__init__()
+        self.max_points = max_points
+        self.conv1 = nn.Conv1d(input_dim, 64, 1)
+        self.conv2 = nn.Conv1d(64, 128, 1)
+        self.conv3 = nn.Conv1d(128, 256, 1)
+        self.conv4 = nn.Conv1d(256, 512, 1)
+        self.conv5 = nn.Conv1d(512, 1024, 1)
+        self.conv6 = nn.Conv1d(1024, 2048, 1)
+        self.fc1 = nn.Linear(2048, 1024)
+        self.fc2 = nn.Linear(1024, 512)
+        self.fc3 = nn.Linear(512, 256)
+        self.fc4 = nn.Linear(256, 128)
+        self.fc5 = nn.Linear(128, 64)
+        self.fc6 = nn.Linear(64, 1)
+        self.bn1 = nn.BatchNorm1d(64)
+        self.bn2 = nn.BatchNorm1d(128)
+        self.bn3 = nn.BatchNorm1d(256)
+        self.bn4 = nn.BatchNorm1d(512)
+        self.bn5 = nn.BatchNorm1d(1024)
+        self.bn6 = nn.BatchNorm1d(2048)
+        self.dropout1 = nn.Dropout(0.3)
+        self.dropout2 = nn.Dropout(0.4)
+        self.dropout3 = nn.Dropout(0.5)
+        self.dropout4 = nn.Dropout(0.4)
+        self.dropout5 = nn.Dropout(0.3)
+
+    def forward(self, x):
+        # x: (B, 6, max_points)
+        x1 = F.relu(self.bn1(self.conv1(x)))
+        x2 = F.relu(self.bn2(self.conv2(x1)))
+        x3 = F.relu(self.bn3(self.conv3(x2)))
+        x4 = F.relu(self.bn4(self.conv4(x3)))
+        x5 = F.relu(self.bn5(self.conv5(x4)))
+        x6 = F.relu(self.bn6(self.conv6(x5)))
+        g = torch.max(x6, 2)[0]
+        x = F.relu(self.fc1(g));   x = self.dropout1(x)
+        x = F.relu(self.fc2(x));   x = self.dropout2(x)
+        x = F.relu(self.fc3(x));   x = self.dropout3(x)
+        x = F.relu(self.fc4(x));   x = self.dropout4(x)
+        x = F.relu(self.fc5(x));   x = self.dropout5(x)
+        return self.fc6(x)         # (B, 1) logits
+
+
+def load_pnet_class(model_path, device=None):
+    """Load a trained ClassificationPointNet checkpoint."""
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = ClassificationPointNet(input_dim=6, max_points=MAX_PATCH_POINTS)
+    ckpt = torch.load(model_path, map_location=device, weights_only=False)
+    model.load_state_dict(ckpt['model_state_dict'])
+    model.to(device).eval()
+    return model
+
+
+def colmap_points_xyz_rgb(colmap_rec):
+    """Return (xyz, rgb_normalized_0_1) for all COLMAP points."""
+    xyz_list, rgb_list = [], []
+    for _, p3D in colmap_rec.points3D.items():
+        xyz_list.append(p3D.xyz)
+        rgb_list.append(p3D.color / 255.0)
+    if not xyz_list:
+        return np.empty((0, 3)), np.empty((0, 3))
+    return np.array(xyz_list), np.array(rgb_list)
+
+
+def build_edge_patch_6d(u_xyz, v_xyz, colmap_xyz, colmap_rgb):
+    """6D cylinder patch around edge (u, v). None if too sparse."""
+    line = v_xyz - u_xyz
+    L = float(np.linalg.norm(line))
+    if L < 1e-6:
+        return None
+    direction = line / L
+    ext_start = u_xyz - CYL_EXT * direction
+    ext_L = L + 2 * CYL_EXT
+    rel = colmap_xyz - ext_start[np.newaxis, :]
+    proj = rel @ direction
+    in_bounds = (proj >= 0) & (proj <= ext_L)
+    closest = ext_start[np.newaxis, :] + proj[:, np.newaxis] * direction[np.newaxis, :]
+    perp = np.linalg.norm(colmap_xyz - closest, axis=1)
+    in_cyl = in_bounds & (perp <= CYL_RADIUS)
+    if int(in_cyl.sum()) <= 10:
+        return None
+    midpoint = (u_xyz + v_xyz) / 2
+    pts_centered = colmap_xyz[in_cyl] - midpoint
+    rgb_signed = colmap_rgb[in_cyl] * 2.0 - 1.0
+    return np.hstack([pts_centered, rgb_signed])
+
+
+def _pad_or_sample_patch(patch_6d, max_pts=MAX_PATCH_POINTS, rng=None):
+    if rng is None:
+        rng = np.random
+    n = patch_6d.shape[0]
+    if n >= max_pts:
+        idx = rng.choice(n, max_pts, replace=False)
+        return patch_6d[idx]
+    out = np.zeros((max_pts, 6), dtype=np.float32)
+    out[:n] = patch_6d
+    return out
+
+
+def score_edges_batched(model, device, vertices, edges, colmap_xyz, colmap_rgb,
+                         rng=None, batch_size=32):
+    """Score every edge in `edges` (returns list of floats, or None per failed patch)."""
+    if rng is None:
+        rng = np.random.RandomState(0)
+    n = len(edges)
+    scores = [None] * n
+    if n == 0 or len(vertices) == 0 or len(colmap_xyz) == 0:
+        return scores
+
+    patches, indices = [], []
+    for i, (u, v) in enumerate(edges):
+        u_xyz = vertices[int(u)]
+        v_xyz = vertices[int(v)]
+        raw = build_edge_patch_6d(u_xyz, v_xyz, colmap_xyz, colmap_rgb)
+        if raw is None:
+            continue
+        patches.append(_pad_or_sample_patch(raw, rng=rng).astype(np.float32))
+        indices.append(i)
+
+    if not patches:
+        return scores
+
+    batch = np.stack(patches, axis=0).transpose(0, 2, 1)  # (N, 6, 1024)
+    with torch.no_grad():
+        for start in range(0, len(batch), batch_size):
+            end = min(start + batch_size, len(batch))
+            x = torch.from_numpy(batch[start:end]).to(device)
+            logits = model(x).squeeze(-1)
+            probs = torch.sigmoid(logits).cpu().numpy().reshape(-1)
+            for j, p in enumerate(probs):
+                scores[indices[start + j]] = float(p)
+    return scores
+
+
+def augment_hybrid_with_filtered_hc(h_v, h_e, user_v, user_e, scores,
+                                     thresh=AUGMENT_THRESHOLD,
+                                     dedup_radius=AUGMENT_DEDUP_RADIUS):
+    """Add handcrafted edges scoring above thresh into the hybrid (v, e)."""
+    h_v = np.asarray(h_v, dtype=np.float32)
+    user_v = np.asarray(user_v, dtype=np.float32)
+    h_e_list = [(int(a), int(b)) for a, b in h_e]
+    if scores is None or len(user_e) == 0 or len(user_v) == 0:
+        return h_v, h_e_list
+
+    kept_pairs = [
+        (int(u), int(v)) for (u, v), s in zip(user_e, scores)
+        if s is not None and s > thresh
+    ]
+    if not kept_pairs:
+        return h_v, h_e_list
+
+    needed_hc_idx = sorted({u for u, _ in kept_pairs} | {v for _, v in kept_pairs})
+    new_v_list, user_to_combined = [], {}
+    for u_idx in needed_hc_idx:
+        pos = user_v[u_idx]
+        if len(h_v) > 0:
+            d = np.linalg.norm(h_v - pos, axis=1)
+            best = int(np.argmin(d))
+            if d[best] <= dedup_radius:
+                user_to_combined[u_idx] = best
+                continue
+        user_to_combined[u_idx] = len(h_v) + len(new_v_list)
+        new_v_list.append(pos)
+
+    if new_v_list:
+        combined_v = np.concatenate([h_v, np.stack(new_v_list, axis=0)], axis=0)
+    else:
+        combined_v = h_v
+
+    existing = {(min(a, b), max(a, b)) for a, b in h_e_list}
+    new_edges = []
+    for u, v in kept_pairs:
+        a, b = user_to_combined[u], user_to_combined[v]
+        if a == b:
+            continue
+        key = (min(a, b), max(a, b))
+        if key in existing:
+            continue
+        existing.add(key)
+        new_edges.append((a, b))
+    return combined_v, h_e_list + new_edges

experiments/ptv3/README.md

@@ -0,0 +1,32 @@
+# Point Transformer V3 encoder (negative result)
+
+This is the Point Transformer V3 encoder experiment from the write-up. We
+replaced the Perceiver encoder with a Point Transformer V3 encoder (per-point
+features, no latent bottleneck) to test whether a stronger encoder would scale
+past the 8k Perceiver.
+
+**Outcome.** Trained from scratch at 8k for 200k steps, it plateaued at ~0.323
+local HSS, below the curriculum-trained Perceiver (~0.357), and ran at roughly
+6 s/sample on an A5000 — over the two-hour T4 evaluation budget. It was never
+submitted. This folder is provided as confirmation that the experiment was run.
+
+## Contents
+
+```
+checkpoint_ptv3_8k.pt     trained weights (200k steps at 8k)
+train_args.json           training configuration
+ptv3_code/                the PT v3 encoder (adapted from the Pointcept release)
+                          and the [B,T,*] <-> flat adapter (encoder_wrapper.py)
+model_with_ptv3.py        EdgeDepthSegmentsModel with the arch="ptv3" branch
+                          (drop-in replacement for s23dr_2026_example/model.py)
+```
+
+## Notes
+
+- Extra dependencies beyond the main `requirements.txt`:
+  `spconv-cu121`, `torch-scatter`, `addict` (and optionally `flash-attn`).
+- To run it, place `ptv3_code/` as `s23dr_2026_example/ptv3/`, use
+  `model_with_ptv3.py` as `s23dr_2026_example/model.py`, and load the checkpoint
+  with `arch="ptv3"`. Inference must use fp32 (spconv does not tune fp16 kernels
+  reliably on these GPUs).
+- The PT v3 encoder is adapted from the authors' Pointcept release.

experiments/ptv3/checkpoint_ptv3_8k.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:317879f9efd320246dafa65132a55669ec39691bbcbf5ab2f9e24f94767b6ce1
+size 461100495

experiments/ptv3/model_with_ptv3.py

@@ -0,0 +1,698 @@
+"""
+Perceiver-based transformer for 3D roof wireframe prediction.
+
+Architecture overview:
+
+    Input tokens [B, T, D]
+        |
+        v
+    input_proj: Linear -> GELU -> Linear -> LayerNorm   =>  [B, T, hidden]
+        |
+        v
+    Perceiver latent bottleneck (N PerceiverLatentLayers):
+        Learnable latent embeddings [L, hidden] are broadcast to batch.
+        Each layer: cross-attn(latents <- tokens) -> self-attn(latents) -> FFN
+        Output: latents [B, L, hidden]
+        |
+        v
+    Segment decoder (M SegmentDecoderLayers):
+        Learnable query embeddings [S, hidden] are broadcast to batch.
+        Each layer: cross-attn(queries <- latents) -> self-attn(queries) -> FFN
+        Output: queries [B, S, hidden]
+        |
+        v
+    segment_head: Linear -> 6D -> (midpoint, half_vector)
+        + query_offsets (learnable per-query bias)
+        endpoints = midpoint +/- half_vector  ->  [B, S, 2, 3]
+"""
+
+import torch
+import torch.nn as nn
+
+from .attention import MultiHeadSDPA, FeedForward
+
+
+# ---------------------------------------------------------------------------
+# Building blocks
+# ---------------------------------------------------------------------------
+
+class AttnResidual(nn.Module):
+    """Pre-norm attention + residual + dropout."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        kv_heads: int | None = None,
+        norm_class=None,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        norm_class = norm_class or nn.LayerNorm
+        self.norm = norm_class(d_model)
+        self.attn = MultiHeadSDPA(d_model, num_heads, kv_heads=kv_heads, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        memory: torch.Tensor,
+        memory_key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        res = x
+        x = self.norm(x)
+        x = self.attn(x, memory, key_padding_mask=memory_key_padding_mask)
+        return res + self.drop(x)
+
+
+class FFNResidual(nn.Module):
+    """Pre-norm feed-forward + residual + dropout."""
+
+    def __init__(
+        self,
+        d_model: int,
+        dim_ff: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        norm_class=None,
+    ):
+        super().__init__()
+        norm_class = norm_class or nn.LayerNorm
+        self.norm = norm_class(d_model)
+        self.ffn = FeedForward(d_model, dim_ff, activation=activation)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        res = x
+        x = self.norm(x)
+        x = self.ffn(x)
+        return res + self.drop(x)
+
+
+# ---------------------------------------------------------------------------
+# Perceiver encoder layer
+# ---------------------------------------------------------------------------
+
+class PerceiverLatentLayer(nn.Module):
+    """Single Perceiver latent layer.
+
+    If use_cross=True:  cross-attn(latents <- points) -> self-attn -> FFN
+    If use_cross=False: self-attn -> FFN  (saves compute in deep stacks)
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_ff: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        kv_heads_cross: int | None = None,
+        kv_heads_self: int | None = None,
+        use_cross: bool = True,
+        norm_class=None,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        self.use_cross = use_cross
+        if use_cross:
+            self.cross = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_self, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)
+
+    def forward(
+        self,
+        latents: torch.Tensor,
+        points: torch.Tensor,
+        points_key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.use_cross:
+            latents = self.cross(latents, points, memory_key_padding_mask=points_key_padding_mask)
+        latents = self.self_attn(latents, latents)
+        latents = self.ffn(latents)
+        return latents
+
+
+# ---------------------------------------------------------------------------
+# Segment decoder layer
+# ---------------------------------------------------------------------------
+
+class SegmentDecoderLayer(nn.Module):
+    """Single segment decoder layer.
+
+    cross-attn(queries <- latents) -> [cross-attn(queries <- inputs)] -> self-attn(queries) -> FFN
+
+    If input_xattn=True, adds a second cross-attention that attends directly
+    to the projected input tokens (bypassing the latent bottleneck). This gives
+    queries access to fine-grained point-level detail for vertex precision.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_ff: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        kv_heads_cross: int | None = None,
+        kv_heads_self: int | None = None,
+        norm_class=None,
+        input_xattn: bool = False,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        self.cross = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.input_xattn = input_xattn
+        if input_xattn:
+            self.cross_input = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_self, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)
+
+    def forward(
+        self,
+        queries: torch.Tensor,
+        latents: torch.Tensor,
+        src: torch.Tensor | None = None,
+        src_key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        queries = self.cross(queries, latents)
+        if self.input_xattn and src is not None:
+            queries = self.cross_input(queries, src, memory_key_padding_mask=src_key_padding_mask)
+        queries = self.self_attn(queries, queries)
+        queries = self.ffn(queries)
+        return queries
+
+
+# ---------------------------------------------------------------------------
+# Full model
+# ---------------------------------------------------------------------------
+
+class TokenTransformerSegments(nn.Module):
+    """Perceiver transformer that predicts 3D roof wireframe segments.
+
+    Takes point-cloud tokens and outputs segment endpoints as [B, S, 2, 3]
+    where S is the number of segments and each segment has two 3D endpoints.
+
+    Args:
+        segments:       Number of predicted segments (S).
+        in_dim:         Dimensionality of input tokens.
+        hidden:         Internal hidden dimension throughout the model.
+        num_heads:      Number of attention heads.
+        kv_heads_cross: Grouped-query heads for cross-attention (None = standard MHA).
+        kv_heads_self:  Grouped-query heads for self-attention (None = standard MHA).
+        dim_feedforward: FFN intermediate dimension.
+        dropout:        Dropout rate applied after attention and FFN.
+        latent_tokens:  Number of learnable latent embeddings (L) in the bottleneck.
+        latent_layers:  Number of PerceiverLatentLayers (N).
+        decoder_layers: Number of SegmentDecoderLayers (M).
+    """
+
+    def __init__(
+        self,
+        segments: int = 32,
+        in_dim: int = 128,
+        hidden: int = 128,
+        num_heads: int = 4,
+        kv_heads_cross: int | None = 2,
+        kv_heads_self: int | None = 0,
+        dim_feedforward: int = 256,
+        dropout: float = 0.01,
+        latent_tokens: int = 64,
+        latent_layers: int = 2,
+        decoder_layers: int = 2,
+        cross_attn_interval: int = 1,
+        norm_class=None,
+        activation: str = "gelu",
+        segment_conf: bool = False,
+        pre_encoder_layers: int = 0,
+        segment_param: str = "midpoint_halfvec",
+        length_floor: float = 0.0,
+        decoder_input_xattn: bool = False,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        self.segments = segments
+        self.out_vertices = segments * 2
+        self.segment_param = segment_param
+        self.decoder_input_xattn = decoder_input_xattn
+        norm_class = norm_class or nn.LayerNorm
+
+        # Treat 0 as "use standard MHA"
+        if kv_heads_cross is not None and kv_heads_cross <= 0:
+            kv_heads_cross = None
+        if kv_heads_self is not None and kv_heads_self <= 0:
+            kv_heads_self = None
+
+        # -- Input projection --
+        self.input_proj = nn.Sequential(
+            nn.Linear(in_dim, dim_feedforward),
+            nn.GELU(),
+            nn.Linear(dim_feedforward, hidden),
+            norm_class(hidden),
+        )
+
+        # -- Optional pre-encoder: self-attention on full token sequence --
+        if pre_encoder_layers > 0:
+            self.pre_encoder = nn.ModuleList([
+                SelfAttentionEncoderLayer(
+                    d_model=hidden,
+                    num_heads=num_heads,
+                    dim_ff=dim_feedforward,
+                    dropout=dropout,
+                    activation=activation,
+                    kv_heads=kv_heads_self,
+                    norm_class=norm_class,
+                    qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+                )
+                for _ in range(pre_encoder_layers)
+            ])
+        else:
+            self.pre_encoder = None
+
+        # -- Perceiver latent bottleneck --
+        self.latent_embed = nn.Embedding(latent_tokens, hidden)
+        N = latent_layers
+        self.latent_layers = nn.ModuleList([
+            PerceiverLatentLayer(
+                d_model=hidden,
+                num_heads=num_heads,
+                dim_ff=dim_feedforward,
+                dropout=dropout,
+                activation=activation,
+                kv_heads_cross=kv_heads_cross,
+                kv_heads_self=kv_heads_self,
+                use_cross=(i == 0) or (i == N - 1) or (i % cross_attn_interval == 0),
+                norm_class=norm_class,
+                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+            )
+            for i in range(N)
+        ])
+
+        # -- Segment decoder --
+        self.query_embed = nn.Embedding(segments, hidden)
+        self.decoder_layers = nn.ModuleList([
+            SegmentDecoderLayer(
+                d_model=hidden,
+                num_heads=num_heads,
+                dim_ff=dim_feedforward,
+                dropout=dropout,
+                activation=activation,
+                kv_heads_cross=kv_heads_cross,
+                kv_heads_self=kv_heads_self,
+                norm_class=norm_class,
+                input_xattn=decoder_input_xattn,
+                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+            )
+            for _ in range(decoder_layers)
+        ])
+
+        # -- Output head --
+        if segment_param == "midpoint_dir_len":
+            self.segment_head = nn.Linear(hidden, 7)  # mid(3) + dir(3) + len(1)
+        else:
+            self.segment_head = nn.Linear(hidden, 6)  # mid(3) + half(3)
+        self.query_offsets = nn.Parameter(torch.zeros(segments, 2, 3))
+
+        nn.init.trunc_normal_(self.segment_head.weight, mean=0.0, std=1e-3)
+        if self.segment_head.bias is not None:
+            nn.init.zeros_(self.segment_head.bias)
+            if segment_param == "midpoint_dir_len":
+                # softplus(0.5) * 0.1 ≈ 0.097 default length in normalized space
+                self.segment_head.bias.data[6] = 0.5
+        nn.init.normal_(self.query_offsets, mean=0.0, std=0.05)
+
+        # -- Optional confidence head --
+        self.segment_conf = segment_conf
+        if segment_conf:
+            self.conf_head = nn.Linear(hidden, 1)
+            nn.init.zeros_(self.conf_head.bias)
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        mask: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor | list]:
+        """
+        Args:
+            tokens: Input point-cloud tokens [B, T, in_dim].
+            mask:   Boolean validity mask [B, T]. True = valid token.
+
+        Returns:
+            Dict with keys:
+                "vertices": [B, S*2, 3] flattened endpoints.
+                "segments": [B, S, 2, 3] segment endpoints.
+                "edges":    Per-batch list of (start, end) index pairs into vertices.
+                "conf":     [B, S] logits (only if segment_conf=True).
+        """
+        B = tokens.shape[0]
+
+        # Project input tokens
+        src = self.input_proj(tokens)  # [B, T, hidden]
+
+        # Padding mask (True where padded) for cross-attention
+        pad_mask = ~mask.bool() if mask is not None else None
+
+        # Optional pre-encoder: self-attention on full token sequence
+        if self.pre_encoder is not None:
+            for layer in self.pre_encoder:
+                src = layer(src, key_padding_mask=pad_mask)
+
+        # Perceiver latent bottleneck
+        latents = self.latent_embed.weight.unsqueeze(0).expand(B, -1, -1)
+        for layer in self.latent_layers:
+            latents = layer(latents, src, points_key_padding_mask=pad_mask)
+
+        # Segment decoder
+        queries = self.query_embed.weight.unsqueeze(0).expand(B, -1, -1)
+        for layer in self.decoder_layers:
+            queries = layer(queries, latents,
+                            src=src if self.decoder_input_xattn else None,
+                            src_key_padding_mask=pad_mask if self.decoder_input_xattn else None)
+
+        # Predict segments -> endpoints
+        if self.segment_param == "midpoint_dir_len":
+            raw = self.segment_head(queries)  # [B, S, 7]
+            mid = raw[:, :, :3] + self.query_offsets[:, 0, :].unsqueeze(0)
+            direction = torch.nn.functional.normalize(raw[:, :, 3:6], dim=-1)
+            length = torch.nn.functional.softplus(raw[:, :, 6:7]) * 0.1
+            half = direction * length * 0.5
+        else:
+            raw = self.segment_head(queries).view(B, self.segments, 2, 3)
+            raw = raw + self.query_offsets.unsqueeze(0)
+            mid, half = raw[:, :, 0], raw[:, :, 1]
+        seg_params = torch.stack([mid - half, mid + half], dim=2)
+
+        vertices = seg_params.reshape(B, self.out_vertices, 3)
+        edges = [[(2 * i, 2 * i + 1) for i in range(self.segments)] for _ in range(B)]
+
+        out = {"vertices": vertices, "segments": seg_params, "edges": edges,
+               "src": src, "pad_mask": pad_mask, "queries": queries}
+        if self.segment_conf:
+            out["conf"] = self.conf_head(queries).squeeze(-1)  # [B, S]
+        return out
+
+
+# ---------------------------------------------------------------------------
+# PT v3 segmenter: same decoder + heads as TokenTransformerSegments but uses
+# PointTransformerV3 (per-point features, no Perceiver bottleneck) as encoder.
+# ---------------------------------------------------------------------------
+
+
+class TokenPTv3Segments(nn.Module):
+    """PT v3 encoder + DETR-style segment decoder.
+
+    Differs from TokenTransformerSegments only in the encoder: instead of an
+    input projection + Perceiver latent bottleneck, this uses PointTransformerV3
+    to produce per-point features [B, T, hidden] (no compression). The decoder,
+    output heads, and segment parameterization are identical.
+
+    The first 3 dims of the input `tokens` tensor are assumed to be xyz_norm
+    (per data.build_tokens / tokenizer order).
+    """
+
+    def __init__(
+        self,
+        segments: int = 64,
+        in_dim: int = 95,
+        hidden: int = 256,
+        num_heads: int = 4,
+        kv_heads_cross: int | None = 2,
+        kv_heads_self: int | None = 2,
+        dim_feedforward: int = 1024,
+        dropout: float = 0.1,
+        decoder_layers: int = 3,
+        norm_class=None,
+        activation: str = "gelu",
+        segment_conf: bool = True,
+        segment_param: str = "midpoint_dir_len",
+        length_floor: float = 0.0,
+        decoder_input_xattn: bool = False,
+        qk_norm: bool = True,
+        qk_norm_type: str = "l2",
+        # PT v3 hyperparams (subset of PTv3Encoder kwargs)
+        ptv3_grid_size: float = 0.005,
+        ptv3_enc_channels: tuple = (32, 64, 128, 256, 256),
+        ptv3_enc_depths: tuple = (2, 2, 2, 6, 2),
+        ptv3_enc_num_head: tuple = (2, 4, 8, 16, 16),
+        ptv3_dec_channels: tuple = (256, 64, 128, 256),
+        ptv3_dec_depths: tuple = (2, 2, 2, 2),
+        ptv3_dec_num_head: tuple = (4, 4, 8, 16),
+        ptv3_stride: tuple = (2, 2, 2, 2),
+        ptv3_enable_flash: bool = True,
+        ptv3_drop_path: float = 0.3,
+    ):
+        super().__init__()
+        from .ptv3.encoder_wrapper import PTv3Encoder
+
+        self.segments = segments
+        self.out_vertices = segments * 2
+        self.segment_param = segment_param
+        self.decoder_input_xattn = decoder_input_xattn
+        norm_class = norm_class or nn.LayerNorm
+        if kv_heads_cross is not None and kv_heads_cross <= 0:
+            kv_heads_cross = None
+        if kv_heads_self is not None and kv_heads_self <= 0:
+            kv_heads_self = None
+
+        # PT v3 encoder. in_channels = full token feature dim. coord is sliced
+        # from the first 3 dims of tokens at forward time.
+        self.ptv3_encoder = PTv3Encoder(
+            in_channels=in_dim,
+            hidden=hidden,
+            grid_size=ptv3_grid_size,
+            enc_channels=ptv3_enc_channels,
+            enc_depths=ptv3_enc_depths,
+            enc_num_head=ptv3_enc_num_head,
+            dec_channels=ptv3_dec_channels,
+            dec_depths=ptv3_dec_depths,
+            dec_num_head=ptv3_dec_num_head,
+            stride=ptv3_stride,
+            enable_flash=ptv3_enable_flash,
+            drop_path=ptv3_drop_path,
+        )
+
+        # DETR-style decoder (same as TokenTransformerSegments)
+        self.query_embed = nn.Embedding(segments, hidden)
+        self.decoder_layers = nn.ModuleList([
+            SegmentDecoderLayer(
+                d_model=hidden,
+                num_heads=num_heads,
+                dim_ff=dim_feedforward,
+                dropout=dropout,
+                activation=activation,
+                kv_heads_cross=kv_heads_cross,
+                kv_heads_self=kv_heads_self,
+                norm_class=norm_class,
+                input_xattn=decoder_input_xattn,
+                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+            )
+            for _ in range(decoder_layers)
+        ])
+
+        # Output heads (identical to TokenTransformerSegments)
+        if segment_param == "midpoint_dir_len":
+            self.segment_head = nn.Linear(hidden, 7)
+        else:
+            self.segment_head = nn.Linear(hidden, 6)
+        self.query_offsets = nn.Parameter(torch.zeros(segments, 2, 3))
+        nn.init.trunc_normal_(self.segment_head.weight, mean=0.0, std=1e-3)
+        if self.segment_head.bias is not None:
+            nn.init.zeros_(self.segment_head.bias)
+            if segment_param == "midpoint_dir_len":
+                self.segment_head.bias.data[6] = 0.5
+        nn.init.normal_(self.query_offsets, mean=0.0, std=0.05)
+
+        self.segment_conf = segment_conf
+        if segment_conf:
+            self.conf_head = nn.Linear(hidden, 1)
+            nn.init.zeros_(self.conf_head.bias)
+
+    def forward(self, tokens: torch.Tensor, mask: torch.Tensor):
+        """
+        Args:
+            tokens: [B, T, in_dim] full per-point features. First 3 dims = xyz_norm.
+            mask:   [B, T] boolean validity.
+
+        Returns same dict as TokenTransformerSegments.forward.
+        """
+        B = tokens.shape[0]
+        coord = tokens[..., :3]                             # [B, T, 3]
+        # PT v3 encoder returns per-point features at original token resolution.
+        src = self.ptv3_encoder(coord, tokens, mask=mask)   # [B, T, hidden]
+
+        pad_mask = ~mask.bool() if mask is not None else None
+
+        # DETR-style decoder cross-attends to per-point features.
+        queries = self.query_embed.weight.unsqueeze(0).expand(B, -1, -1)
+        for layer in self.decoder_layers:
+            queries = layer(queries, src,
+                            src=src if self.decoder_input_xattn else None,
+                            src_key_padding_mask=pad_mask if self.decoder_input_xattn else None)
+
+        # Predict segments -> endpoints (identical to TokenTransformerSegments)
+        if self.segment_param == "midpoint_dir_len":
+            raw = self.segment_head(queries)
+            mid = raw[:, :, :3] + self.query_offsets[:, 0, :].unsqueeze(0)
+            direction = torch.nn.functional.normalize(raw[:, :, 3:6], dim=-1)
+            length = torch.nn.functional.softplus(raw[:, :, 6:7]) * 0.1
+            half = direction * length * 0.5
+        else:
+            raw = self.segment_head(queries).view(B, self.segments, 2, 3)
+            raw = raw + self.query_offsets.unsqueeze(0)
+            mid, half = raw[:, :, 0], raw[:, :, 1]
+        seg_params = torch.stack([mid - half, mid + half], dim=2)
+
+        vertices = seg_params.reshape(B, self.out_vertices, 3)
+        edges = [[(2 * i, 2 * i + 1) for i in range(self.segments)] for _ in range(B)]
+
+        out = {"vertices": vertices, "segments": seg_params, "edges": edges,
+               "src": src, "pad_mask": pad_mask, "queries": queries}
+        if self.segment_conf:
+            out["conf"] = self.conf_head(queries).squeeze(-1)
+        return out
+
+
+# ---------------------------------------------------------------------------
+# Encoder-only layer (self-attention on full token sequence)
+# ---------------------------------------------------------------------------
+
+class SelfAttentionEncoderLayer(nn.Module):
+    """Single self-attention layer: self-attn(tokens) -> FFN."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_ff: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        kv_heads: int | None = None,
+        norm_class=None,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)
+
+    def forward(self, x: torch.Tensor, key_padding_mask: torch.Tensor | None = None) -> torch.Tensor:
+        x = self.self_attn(x, x, memory_key_padding_mask=key_padding_mask)
+        x = self.ffn(x)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# End-to-end model: tokenizer embeddings + perceiver
+# ---------------------------------------------------------------------------
+
+class EdgeDepthSegmentsModel(nn.Module):
+    """Tokenizer embeddings + transformer for 3D roof wireframes.
+
+    Supports two architectures via the `arch` parameter:
+    - "perceiver": Perceiver latent bottleneck (default, O(L*T) attention)
+    - "transformer": Standard self-attention encoder (O(T^2) attention)
+
+    Both share the same decoder, output head, and tokenizer.
+    """
+
+    def __init__(
+        self,
+        seq_cfg,
+        segments: int = 32,
+        hidden: int = 128,
+        num_heads: int = 4,
+        kv_heads_cross: int | None = 2,
+        kv_heads_self: int | None = 0,
+        dim_feedforward: int = 256,
+        dropout: float = 0.1,
+        latent_tokens: int = 64,
+        latent_layers: int = 1,
+        decoder_layers: int = 2,
+        label_emb_dim: int = 16,
+        src_emb_dim: int = 2,
+        behind_emb_dim: int = 8,
+        fourier_seed: int = 0,
+        cross_attn_interval: int = 1,
+        norm_class=None,
+        activation: str = "gelu",
+        segment_conf: bool = False,
+        use_vote_features: bool = False,
+        arch: str = "perceiver",
+        encoder_layers: int = 4,
+        pre_encoder_layers: int = 0,
+        segment_param: str = "midpoint_halfvec",
+        length_floor: float = 0.0,
+        decoder_input_xattn: bool = False,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+        learnable_fourier: bool = False,
+    ):
+        super().__init__()
+        self.seq_cfg = seq_cfg
+
+        from .tokenizer import EdgeDepthSequenceBuilder
+        self.tokenizer = EdgeDepthSequenceBuilder(
+            seq_cfg,
+            label_emb_dim=label_emb_dim,
+            src_emb_dim=src_emb_dim,
+            behind_emb_dim=behind_emb_dim,
+            fourier_seed=fourier_seed,
+            use_vote_features=use_vote_features,
+            learnable_fourier=learnable_fourier,
+        )
+
+        if arch == "transformer":
+            raise ValueError(
+                "arch='transformer' is no longer supported. "
+                "TransformerSegments has been removed; use arch='perceiver'.")
+        elif arch == "ptv3":
+            self.segmenter = TokenPTv3Segments(
+                segments=segments,
+                in_dim=self.tokenizer.out_dim,
+                hidden=hidden,
+                num_heads=num_heads,
+                kv_heads_cross=kv_heads_cross,
+                kv_heads_self=kv_heads_self,
+                dim_feedforward=dim_feedforward,
+                dropout=dropout,
+                decoder_layers=decoder_layers,
+                norm_class=norm_class,
+                activation=activation,
+                segment_conf=segment_conf,
+                segment_param=segment_param,
+                length_floor=length_floor,
+                decoder_input_xattn=decoder_input_xattn,
+                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+            )
+        else:
+            self.segmenter = TokenTransformerSegments(
+                segments=segments,
+                in_dim=self.tokenizer.out_dim,
+                hidden=hidden,
+                num_heads=num_heads,
+                kv_heads_cross=kv_heads_cross,
+                kv_heads_self=kv_heads_self,
+                dim_feedforward=dim_feedforward,
+                dropout=dropout,
+                latent_tokens=latent_tokens,
+                latent_layers=latent_layers,
+                decoder_layers=decoder_layers,
+                cross_attn_interval=cross_attn_interval,
+                norm_class=norm_class,
+                activation=activation,
+                segment_conf=segment_conf,
+                pre_encoder_layers=pre_encoder_layers,
+                segment_param=segment_param,
+                length_floor=length_floor,
+                decoder_input_xattn=decoder_input_xattn,
+                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+            )
+
+    def forward_tokens(self, tokens: torch.Tensor, mask: torch.Tensor):
+        """Run the segmenter on pre-built token tensors."""
+        return self.segmenter(tokens, mask)

experiments/ptv3/ptv3_code/encoder_wrapper.py

@@ -0,0 +1,174 @@
+"""Adapter that lets PointTransformerV3 plug into our existing model pipeline.
+
+Bridges our [B, T, ·] batched format with PT v3's flat [Σ T_i, ·] + batch-indices
+format. Handles voxel deduplication via inverse mapping so output point count
+matches input exactly, even when multiple input points fall in the same voxel.
+
+Drop-in replacement for the Perceiver latent bottleneck:
+  - Input: per-token features [B, T, in_dim] and coord [B, T, 3]
+  - Output: per-point features [B, T, hidden_out]
+  - Output then fed to the existing DETR-style segment decoder (cross-attn over T tokens).
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from .model import PointTransformerV3
+
+
+class PTv3Encoder(nn.Module):
+    """Wrap PointTransformerV3 in a [B, T, ·] interface."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden: int = 256,
+        grid_size: float = 0.005,
+        enc_channels: tuple = (32, 64, 128, 256, 256),
+        enc_depths: tuple = (2, 2, 2, 6, 2),
+        enc_num_head: tuple = (2, 4, 8, 16, 16),
+        enc_patch_size: tuple = (1024, 1024, 1024, 1024, 1024),
+        dec_channels: tuple = (256, 64, 128, 256),
+        dec_depths: tuple = (2, 2, 2, 2),
+        dec_num_head: tuple = (4, 4, 8, 16),
+        dec_patch_size: tuple = (1024, 1024, 1024, 1024),
+        stride: tuple = (2, 2, 2, 2),
+        order: tuple = ("z", "hilbert"),
+        enable_flash: bool = False,
+        shuffle_orders: bool = True,
+        drop_path: float = 0.3,
+    ):
+        super().__init__()
+        if dec_channels[0] != hidden:
+            raise ValueError(
+                f"dec_channels[0]={dec_channels[0]} must equal hidden={hidden} "
+                f"so PT v3 output dim matches the rest of the model."
+            )
+        self.hidden = hidden
+        self.grid_size = grid_size
+
+        # Gracefully fall back to standard attention if flash_attn isn't installed.
+        # Same weights work in both modes (mathematically equivalent).
+        if enable_flash:
+            try:
+                import flash_attn  # noqa: F401
+            except ImportError:
+                enable_flash = False
+
+        self.ptv3 = PointTransformerV3(
+            in_channels=in_channels,
+            order=order,
+            stride=stride,
+            enc_depths=enc_depths,
+            enc_channels=enc_channels,
+            enc_num_head=enc_num_head,
+            enc_patch_size=enc_patch_size,
+            dec_depths=dec_depths,
+            dec_channels=list(dec_channels),
+            dec_num_head=dec_num_head,
+            dec_patch_size=dec_patch_size,
+            drop_path=drop_path,
+            enable_flash=enable_flash,
+            shuffle_orders=shuffle_orders,
+            cls_mode=False,
+        )
+
+    def forward(self, coord: torch.Tensor, feat: torch.Tensor,
+                mask: torch.Tensor | None = None) -> torch.Tensor:
+        """
+        Args:
+            coord: [B, T, 3] xyz (normalized to roughly [-1, 1])
+            feat:  [B, T, in_channels] per-token features
+            mask:  [B, T] bool; True=valid. If None, all valid.
+
+        Returns:
+            [B, T, hidden] per-point features. Invalid positions are zeroed.
+        """
+        from addict import Dict
+
+        B, T, _ = coord.shape
+        device = coord.device
+
+        if mask is None:
+            valid = torch.ones(B, T, dtype=torch.bool, device=device)
+        else:
+            valid = mask.bool()
+
+        flat_mask = valid.reshape(-1)                          # [B*T]
+        coord_flat = coord.reshape(-1, 3)[flat_mask]           # [N_valid, 3]
+        feat_flat = feat.reshape(-1, feat.shape[-1])[flat_mask]  # [N_valid, in_channels]
+        batch_per_point = torch.arange(B, device=device).repeat_interleave(T)
+        batch = batch_per_point[flat_mask]                     # [N_valid]
+        N_valid = coord_flat.shape[0]
+
+        # Deduplicate (batch, voxel) cells. Multiple input points falling in the
+        # same voxel get merged into one PT v3 token; we map the output back to
+        # all original points via the inverse mapping.
+        coord_min = coord_flat.min(dim=0).values
+        grid_coord = torch.div(
+            coord_flat - coord_min, self.grid_size, rounding_mode="trunc"
+        ).long()                                               # [N_valid, 3]
+        gmax = grid_coord.max(dim=0).values + 1                # [3]
+        stride_y = gmax[2].item()
+        stride_x = gmax[1].item() * stride_y
+        stride_b = gmax[0].item() * stride_x
+
+        voxel_id = (
+            batch * stride_b
+            + grid_coord[:, 0] * stride_x
+            + grid_coord[:, 1] * stride_y
+            + grid_coord[:, 2]
+        )
+        unique_ids, inverse_idx = torch.unique(voxel_id, return_inverse=True)
+        N_unique = unique_ids.shape[0]
+
+        # For each unique cell, pick its first appearance to define the cell's coord/feat.
+        # (PT v3 will internally rebuild structure, this is just initialization.)
+        perm = torch.argsort(inverse_idx, stable=True)
+        first_in_unique = torch.empty(N_unique, dtype=torch.long, device=device)
+        # scatter: for each i in perm in order, last write wins; but with stable sort
+        # the first occurrence has lowest position index, so we need to write in reverse.
+        first_in_unique.scatter_(
+            0, inverse_idx[perm].flip(0), perm.flip(0)
+        )
+
+        coord_u = coord_flat[first_in_unique]                  # [N_unique, 3]
+        feat_u = feat_flat[first_in_unique].contiguous()       # [N_unique, in_channels]
+        batch_u = batch[first_in_unique].contiguous()          # [N_unique]
+        # Sort by batch (PT v3 expects ascending batch)
+        sort_perm = torch.argsort(batch_u, stable=True)
+        coord_u = coord_u[sort_perm].contiguous()
+        feat_u = feat_u[sort_perm].contiguous()
+        batch_u = batch_u[sort_perm].contiguous()
+        # We need inverse mapping that says: for each original point, where is its unique-cell after sort?
+        inv_sort = torch.empty_like(sort_perm)
+        inv_sort[sort_perm] = torch.arange(N_unique, device=device)
+        sorted_unique_idx_for_orig = inv_sort[inverse_idx]     # [N_valid]
+
+        # Build the Point dict and run PT v3.
+        data_dict = Dict(
+            coord=coord_u,
+            feat=feat_u,
+            batch=batch_u,
+            grid_size=self.grid_size,
+        )
+        out_point = self.ptv3(data_dict)
+        out_feat_unique = out_point.feat                       # [N_unique_out, hidden]
+
+        # PT v3 with cls_mode=False restores original (unique-cell) point count, so
+        # N_unique_out == N_unique.
+        if out_feat_unique.shape[0] != N_unique:
+            raise RuntimeError(
+                f"PT v3 output count {out_feat_unique.shape[0]} != input unique count {N_unique}. "
+                f"Did you set cls_mode=False?"
+            )
+
+        # Map unique-cell features back to all original valid points (duplicates share features).
+        per_point_feat = out_feat_unique[sorted_unique_idx_for_orig]  # [N_valid, hidden]
+
+        # Scatter back into [B*T, hidden] then reshape to [B, T, hidden]
+        out = torch.zeros(B * T, self.hidden,
+                          device=device, dtype=per_point_feat.dtype)
+        out[flat_mask] = per_point_feat
+        return out.reshape(B, T, self.hidden)

experiments/ptv3/ptv3_code/model.py

@@ -0,0 +1,982 @@
+"""
+Point Transformer - V3 Mode1
+Pointcept detached version
+
+Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com)
+Please cite our work if the code is helpful to you.
+"""
+
+import sys
+from functools import partial
+from addict import Dict
+import math
+import torch
+import torch.nn as nn
+import spconv.pytorch as spconv
+import torch_scatter
+from timm.models.layers import DropPath
+from collections import OrderedDict
+
+try:
+    import flash_attn
+except ImportError:
+    flash_attn = None
+
+from .serialization import encode
+
+
+@torch.inference_mode()
+def offset2bincount(offset):
+    return torch.diff(
+        offset, prepend=torch.tensor([0], device=offset.device, dtype=torch.long)
+    )
+
+
+@torch.inference_mode()
+def offset2batch(offset):
+    bincount = offset2bincount(offset)
+    return torch.arange(
+        len(bincount), device=offset.device, dtype=torch.long
+    ).repeat_interleave(bincount)
+
+
+@torch.inference_mode()
+def batch2offset(batch):
+    return torch.cumsum(batch.bincount(), dim=0).long()
+
+
+class Point(Dict):
+    """
+    Point Structure of Pointcept
+
+    A Point (point cloud) in Pointcept is a dictionary that contains various properties of
+    a batched point cloud. The property with the following names have a specific definition
+    as follows:
+
+    - "coord": original coordinate of point cloud;
+    - "grid_coord": grid coordinate for specific grid size (related to GridSampling);
+    Point also support the following optional attributes:
+    - "offset": if not exist, initialized as batch size is 1;
+    - "batch": if not exist, initialized as batch size is 1;
+    - "feat": feature of point cloud, default input of model;
+    - "grid_size": Grid size of point cloud (related to GridSampling);
+    (related to Serialization)
+    - "serialized_depth": depth of serialization, 2 ** depth * grid_size describe the maximum of point cloud range;
+    - "serialized_code": a list of serialization codes;
+    - "serialized_order": a list of serialization order determined by code;
+    - "serialized_inverse": a list of inverse mapping determined by code;
+    (related to Sparsify: SpConv)
+    - "sparse_shape": Sparse shape for Sparse Conv Tensor;
+    - "sparse_conv_feat": SparseConvTensor init with information provide by Point;
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # If one of "offset" or "batch" do not exist, generate by the existing one
+        if "batch" not in self.keys() and "offset" in self.keys():
+            self["batch"] = offset2batch(self.offset)
+        elif "offset" not in self.keys() and "batch" in self.keys():
+            self["offset"] = batch2offset(self.batch)
+
+    def serialization(self, order="z", depth=None, shuffle_orders=False):
+        """
+        Point Cloud Serialization
+
+        relay on ["grid_coord" or "coord" + "grid_size", "batch", "feat"]
+        """
+        assert "batch" in self.keys()
+        if "grid_coord" not in self.keys():
+            # if you don't want to operate GridSampling in data augmentation,
+            # please add the following augmentation into your pipline:
+            # dict(type="Copy", keys_dict={"grid_size": 0.01}),
+            # (adjust `grid_size` to what your want)
+            assert {"grid_size", "coord"}.issubset(self.keys())
+            self["grid_coord"] = torch.div(
+                self.coord - self.coord.min(0)[0], self.grid_size, rounding_mode="trunc"
+            ).int()
+
+        if depth is None:
+            # Adaptive measure the depth of serialization cube (length = 2 ^ depth)
+            depth = int(self.grid_coord.max()).bit_length()
+        self["serialized_depth"] = depth
+        # Maximum bit length for serialization code is 63 (int64)
+        assert depth * 3 + len(self.offset).bit_length() <= 63
+        # Here we follow OCNN and set the depth limitation to 16 (48bit) for the point position.
+        # Although depth is limited to less than 16, we can encode a 655.36^3 (2^16 * 0.01) meter^3
+        # cube with a grid size of 0.01 meter. We consider it is enough for the current stage.
+        # We can unlock the limitation by optimizing the z-order encoding function if necessary.
+        assert depth <= 16
+
+        # The serialization codes are arranged as following structures:
+        # [Order1 ([n]),
+        #  Order2 ([n]),
+        #   ...
+        #  OrderN ([n])] (k, n)
+        code = [
+            encode(self.grid_coord, self.batch, depth, order=order_) for order_ in order
+        ]
+        code = torch.stack(code)
+        order = torch.argsort(code)
+        inverse = torch.zeros_like(order).scatter_(
+            dim=1,
+            index=order,
+            src=torch.arange(0, code.shape[1], device=order.device).repeat(
+                code.shape[0], 1
+            ),
+        )
+
+        if shuffle_orders:
+            perm = torch.randperm(code.shape[0])
+            code = code[perm]
+            order = order[perm]
+            inverse = inverse[perm]
+
+        self["serialized_code"] = code
+        self["serialized_order"] = order
+        self["serialized_inverse"] = inverse
+
+    def sparsify(self, pad=96):
+        """
+        Point Cloud Serialization
+
+        Point cloud is sparse, here we use "sparsify" to specifically refer to
+        preparing "spconv.SparseConvTensor" for SpConv.
+
+        relay on ["grid_coord" or "coord" + "grid_size", "batch", "feat"]
+
+        pad: padding sparse for sparse shape.
+        """
+        assert {"feat", "batch"}.issubset(self.keys())
+        if "grid_coord" not in self.keys():
+            # if you don't want to operate GridSampling in data augmentation,
+            # please add the following augmentation into your pipline:
+            # dict(type="Copy", keys_dict={"grid_size": 0.01}),
+            # (adjust `grid_size` to what your want)
+            assert {"grid_size", "coord"}.issubset(self.keys())
+            self["grid_coord"] = torch.div(
+                self.coord - self.coord.min(0)[0], self.grid_size, rounding_mode="trunc"
+            ).int()
+        if "sparse_shape" in self.keys():
+            sparse_shape = self.sparse_shape
+        else:
+            sparse_shape = torch.add(
+                torch.max(self.grid_coord, dim=0).values, pad
+            ).tolist()
+        sparse_conv_feat = spconv.SparseConvTensor(
+            features=self.feat,
+            indices=torch.cat(
+                [self.batch.unsqueeze(-1).int(), self.grid_coord.int()], dim=1
+            ).contiguous(),
+            spatial_shape=sparse_shape,
+            batch_size=self.batch[-1].tolist() + 1,
+        )
+        self["sparse_shape"] = sparse_shape
+        self["sparse_conv_feat"] = sparse_conv_feat
+
+
+class PointModule(nn.Module):
+    r"""PointModule
+    placeholder, all module subclass from this will take Point in PointSequential.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+
+class PointSequential(PointModule):
+    r"""A sequential container.
+    Modules will be added to it in the order they are passed in the constructor.
+    Alternatively, an ordered dict of modules can also be passed in.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        if len(args) == 1 and isinstance(args[0], OrderedDict):
+            for key, module in args[0].items():
+                self.add_module(key, module)
+        else:
+            for idx, module in enumerate(args):
+                self.add_module(str(idx), module)
+        for name, module in kwargs.items():
+            if sys.version_info < (3, 6):
+                raise ValueError("kwargs only supported in py36+")
+            if name in self._modules:
+                raise ValueError("name exists.")
+            self.add_module(name, module)
+
+    def __getitem__(self, idx):
+        if not (-len(self) <= idx < len(self)):
+            raise IndexError("index {} is out of range".format(idx))
+        if idx < 0:
+            idx += len(self)
+        it = iter(self._modules.values())
+        for i in range(idx):
+            next(it)
+        return next(it)
+
+    def __len__(self):
+        return len(self._modules)
+
+    def add(self, module, name=None):
+        if name is None:
+            name = str(len(self._modules))
+            if name in self._modules:
+                raise KeyError("name exists")
+        self.add_module(name, module)
+
+    def forward(self, input):
+        for k, module in self._modules.items():
+            # Point module
+            if isinstance(module, PointModule):
+                input = module(input)
+            # Spconv module
+            elif spconv.modules.is_spconv_module(module):
+                if isinstance(input, Point):
+                    input.sparse_conv_feat = module(input.sparse_conv_feat)
+                    input.feat = input.sparse_conv_feat.features
+                else:
+                    input = module(input)
+            # PyTorch module
+            else:
+                if isinstance(input, Point):
+                    input.feat = module(input.feat)
+                    if "sparse_conv_feat" in input.keys():
+                        input.sparse_conv_feat = input.sparse_conv_feat.replace_feature(
+                            input.feat
+                        )
+                elif isinstance(input, spconv.SparseConvTensor):
+                    if input.indices.shape[0] != 0:
+                        input = input.replace_feature(module(input.features))
+                else:
+                    input = module(input)
+        return input
+
+
+class PDNorm(PointModule):
+    def __init__(
+        self,
+        num_features,
+        norm_layer,
+        context_channels=256,
+        conditions=("ScanNet", "S3DIS", "Structured3D"),
+        decouple=True,
+        adaptive=False,
+    ):
+        super().__init__()
+        self.conditions = conditions
+        self.decouple = decouple
+        self.adaptive = adaptive
+        if self.decouple:
+            self.norm = nn.ModuleList([norm_layer(num_features) for _ in conditions])
+        else:
+            self.norm = norm_layer
+        if self.adaptive:
+            self.modulation = nn.Sequential(
+                nn.SiLU(), nn.Linear(context_channels, 2 * num_features, bias=True)
+            )
+
+    def forward(self, point):
+        assert {"feat", "condition"}.issubset(point.keys())
+        if isinstance(point.condition, str):
+            condition = point.condition
+        else:
+            condition = point.condition[0]
+        if self.decouple:
+            assert condition in self.conditions
+            norm = self.norm[self.conditions.index(condition)]
+        else:
+            norm = self.norm
+        point.feat = norm(point.feat)
+        if self.adaptive:
+            assert "context" in point.keys()
+            shift, scale = self.modulation(point.context).chunk(2, dim=1)
+            point.feat = point.feat * (1.0 + scale) + shift
+        return point
+
+
+class RPE(torch.nn.Module):
+    def __init__(self, patch_size, num_heads):
+        super().__init__()
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.pos_bnd = int((4 * patch_size) ** (1 / 3) * 2)
+        self.rpe_num = 2 * self.pos_bnd + 1
+        self.rpe_table = torch.nn.Parameter(torch.zeros(3 * self.rpe_num, num_heads))
+        torch.nn.init.trunc_normal_(self.rpe_table, std=0.02)
+
+    def forward(self, coord):
+        idx = (
+            coord.clamp(-self.pos_bnd, self.pos_bnd)  # clamp into bnd
+            + self.pos_bnd  # relative position to positive index
+            + torch.arange(3, device=coord.device) * self.rpe_num  # x, y, z stride
+        )
+        out = self.rpe_table.index_select(0, idx.reshape(-1))
+        out = out.view(idx.shape + (-1,)).sum(3)
+        out = out.permute(0, 3, 1, 2)  # (N, K, K, H) -> (N, H, K, K)
+        return out
+
+
+class SerializedAttention(PointModule):
+    def __init__(
+        self,
+        channels,
+        num_heads,
+        patch_size,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        order_index=0,
+        enable_rpe=False,
+        enable_flash=True,
+        upcast_attention=True,
+        upcast_softmax=True,
+    ):
+        super().__init__()
+        assert channels % num_heads == 0
+        self.channels = channels
+        self.num_heads = num_heads
+        self.scale = qk_scale or (channels // num_heads) ** -0.5
+        self.order_index = order_index
+        self.upcast_attention = upcast_attention
+        self.upcast_softmax = upcast_softmax
+        self.enable_rpe = enable_rpe
+        self.enable_flash = enable_flash
+        if enable_flash:
+            assert (
+                enable_rpe is False
+            ), "Set enable_rpe to False when enable Flash Attention"
+            assert (
+                upcast_attention is False
+            ), "Set upcast_attention to False when enable Flash Attention"
+            assert (
+                upcast_softmax is False
+            ), "Set upcast_softmax to False when enable Flash Attention"
+            assert flash_attn is not None, "Make sure flash_attn is installed."
+            self.patch_size = patch_size
+            self.attn_drop = attn_drop
+        else:
+            # when disable flash attention, we still don't want to use mask
+            # consequently, patch size will auto set to the
+            # min number of patch_size_max and number of points
+            self.patch_size_max = patch_size
+            self.patch_size = 0
+            self.attn_drop = torch.nn.Dropout(attn_drop)
+
+        self.qkv = torch.nn.Linear(channels, channels * 3, bias=qkv_bias)
+        self.proj = torch.nn.Linear(channels, channels)
+        self.proj_drop = torch.nn.Dropout(proj_drop)
+        self.softmax = torch.nn.Softmax(dim=-1)
+        self.rpe = RPE(patch_size, num_heads) if self.enable_rpe else None
+
+    @torch.no_grad()
+    def get_rel_pos(self, point, order):
+        K = self.patch_size
+        rel_pos_key = f"rel_pos_{self.order_index}"
+        if rel_pos_key not in point.keys():
+            grid_coord = point.grid_coord[order]
+            grid_coord = grid_coord.reshape(-1, K, 3)
+            point[rel_pos_key] = grid_coord.unsqueeze(2) - grid_coord.unsqueeze(1)
+        return point[rel_pos_key]
+
+    @torch.no_grad()
+    def get_padding_and_inverse(self, point):
+        pad_key = "pad"
+        unpad_key = "unpad"
+        cu_seqlens_key = "cu_seqlens_key"
+        if (
+            pad_key not in point.keys()
+            or unpad_key not in point.keys()
+            or cu_seqlens_key not in point.keys()
+        ):
+            offset = point.offset
+            bincount = offset2bincount(offset)
+            bincount_pad = (
+                torch.div(
+                    bincount + self.patch_size - 1,
+                    self.patch_size,
+                    rounding_mode="trunc",
+                )
+                * self.patch_size
+            )
+            # only pad point when num of points larger than patch_size
+            mask_pad = bincount > self.patch_size
+            bincount_pad = ~mask_pad * bincount + mask_pad * bincount_pad
+            _offset = nn.functional.pad(offset, (1, 0))
+            _offset_pad = nn.functional.pad(torch.cumsum(bincount_pad, dim=0), (1, 0))
+            pad = torch.arange(_offset_pad[-1], device=offset.device)
+            unpad = torch.arange(_offset[-1], device=offset.device)
+            cu_seqlens = []
+            for i in range(len(offset)):
+                unpad[_offset[i] : _offset[i + 1]] += _offset_pad[i] - _offset[i]
+                if bincount[i] != bincount_pad[i]:
+                    pad[
+                        _offset_pad[i + 1]
+                        - self.patch_size
+                        + (bincount[i] % self.patch_size) : _offset_pad[i + 1]
+                    ] = pad[
+                        _offset_pad[i + 1]
+                        - 2 * self.patch_size
+                        + (bincount[i] % self.patch_size) : _offset_pad[i + 1]
+                        - self.patch_size
+                    ]
+                pad[_offset_pad[i] : _offset_pad[i + 1]] -= _offset_pad[i] - _offset[i]
+                cu_seqlens.append(
+                    torch.arange(
+                        _offset_pad[i],
+                        _offset_pad[i + 1],
+                        step=self.patch_size,
+                        dtype=torch.int32,
+                        device=offset.device,
+                    )
+                )
+            point[pad_key] = pad
+            point[unpad_key] = unpad
+            point[cu_seqlens_key] = nn.functional.pad(
+                torch.concat(cu_seqlens), (0, 1), value=_offset_pad[-1]
+            )
+        return point[pad_key], point[unpad_key], point[cu_seqlens_key]
+
+    def forward(self, point):
+        if not self.enable_flash:
+            self.patch_size = min(
+                offset2bincount(point.offset).min().tolist(), self.patch_size_max
+            )
+
+        H = self.num_heads
+        K = self.patch_size
+        C = self.channels
+
+        pad, unpad, cu_seqlens = self.get_padding_and_inverse(point)
+
+        order = point.serialized_order[self.order_index][pad]
+        inverse = unpad[point.serialized_inverse[self.order_index]]
+
+        # padding and reshape feat and batch for serialized point patch
+        qkv = self.qkv(point.feat)[order]
+
+        if not self.enable_flash:
+            # encode and reshape qkv: (N', K, 3, H, C') => (3, N', H, K, C')
+            q, k, v = (
+                qkv.reshape(-1, K, 3, H, C // H).permute(2, 0, 3, 1, 4).unbind(dim=0)
+            )
+            # attn
+            if self.upcast_attention:
+                q = q.float()
+                k = k.float()
+            attn = (q * self.scale) @ k.transpose(-2, -1)  # (N', H, K, K)
+            if self.enable_rpe:
+                attn = attn + self.rpe(self.get_rel_pos(point, order))
+            if self.upcast_softmax:
+                attn = attn.float()
+            attn = self.softmax(attn)
+            attn = self.attn_drop(attn).to(qkv.dtype)
+            feat = (attn @ v).transpose(1, 2).reshape(-1, C)
+        else:
+            feat = flash_attn.flash_attn_varlen_qkvpacked_func(
+                qkv.half().reshape(-1, 3, H, C // H),
+                cu_seqlens,
+                max_seqlen=self.patch_size,
+                dropout_p=self.attn_drop if self.training else 0,
+                softmax_scale=self.scale,
+            ).reshape(-1, C)
+            feat = feat.to(qkv.dtype)
+        feat = feat[inverse]
+
+        # ffn
+        feat = self.proj(feat)
+        feat = self.proj_drop(feat)
+        point.feat = feat
+        return point
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels=None,
+        out_channels=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_channels = out_channels or in_channels
+        hidden_channels = hidden_channels or in_channels
+        self.fc1 = nn.Linear(in_channels, hidden_channels)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_channels, out_channels)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Block(PointModule):
+    def __init__(
+        self,
+        channels,
+        num_heads,
+        patch_size=48,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        drop_path=0.0,
+        norm_layer=nn.LayerNorm,
+        act_layer=nn.GELU,
+        pre_norm=True,
+        order_index=0,
+        cpe_indice_key=None,
+        enable_rpe=False,
+        enable_flash=True,
+        upcast_attention=True,
+        upcast_softmax=True,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.pre_norm = pre_norm
+
+        self.cpe = PointSequential(
+            spconv.SubMConv3d(
+                channels,
+                channels,
+                kernel_size=3,
+                bias=True,
+                indice_key=cpe_indice_key,
+            ),
+            nn.Linear(channels, channels),
+            norm_layer(channels),
+        )
+
+        self.norm1 = PointSequential(norm_layer(channels))
+        self.attn = SerializedAttention(
+            channels=channels,
+            patch_size=patch_size,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            order_index=order_index,
+            enable_rpe=enable_rpe,
+            enable_flash=enable_flash,
+            upcast_attention=upcast_attention,
+            upcast_softmax=upcast_softmax,
+        )
+        self.norm2 = PointSequential(norm_layer(channels))
+        self.mlp = PointSequential(
+            MLP(
+                in_channels=channels,
+                hidden_channels=int(channels * mlp_ratio),
+                out_channels=channels,
+                act_layer=act_layer,
+                drop=proj_drop,
+            )
+        )
+        self.drop_path = PointSequential(
+            DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        )
+
+    def forward(self, point: Point):
+        shortcut = point.feat
+        point = self.cpe(point)
+        point.feat = shortcut + point.feat
+        shortcut = point.feat
+        if self.pre_norm:
+            point = self.norm1(point)
+        point = self.drop_path(self.attn(point))
+        point.feat = shortcut + point.feat
+        if not self.pre_norm:
+            point = self.norm1(point)
+
+        shortcut = point.feat
+        if self.pre_norm:
+            point = self.norm2(point)
+        point = self.drop_path(self.mlp(point))
+        point.feat = shortcut + point.feat
+        if not self.pre_norm:
+            point = self.norm2(point)
+        point.sparse_conv_feat = point.sparse_conv_feat.replace_feature(point.feat)
+        return point
+
+
+class SerializedPooling(PointModule):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride=2,
+        norm_layer=None,
+        act_layer=None,
+        reduce="max",
+        shuffle_orders=True,
+        traceable=True,  # record parent and cluster
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        assert stride == 2 ** (math.ceil(stride) - 1).bit_length()  # 2, 4, 8
+        # TODO: add support to grid pool (any stride)
+        self.stride = stride
+        assert reduce in ["sum", "mean", "min", "max"]
+        self.reduce = reduce
+        self.shuffle_orders = shuffle_orders
+        self.traceable = traceable
+
+        self.proj = nn.Linear(in_channels, out_channels)
+        if norm_layer is not None:
+            self.norm = PointSequential(norm_layer(out_channels))
+        if act_layer is not None:
+            self.act = PointSequential(act_layer())
+
+    def forward(self, point: Point):
+        pooling_depth = (math.ceil(self.stride) - 1).bit_length()
+        if pooling_depth > point.serialized_depth:
+            pooling_depth = 0
+        assert {
+            "serialized_code",
+            "serialized_order",
+            "serialized_inverse",
+            "serialized_depth",
+        }.issubset(
+            point.keys()
+        ), "Run point.serialization() point cloud before SerializedPooling"
+
+        code = point.serialized_code >> pooling_depth * 3
+        code_, cluster, counts = torch.unique(
+            code[0],
+            sorted=True,
+            return_inverse=True,
+            return_counts=True,
+        )
+        # indices of point sorted by cluster, for torch_scatter.segment_csr
+        _, indices = torch.sort(cluster)
+        # index pointer for sorted point, for torch_scatter.segment_csr
+        idx_ptr = torch.cat([counts.new_zeros(1), torch.cumsum(counts, dim=0)])
+        # head_indices of each cluster, for reduce attr e.g. code, batch
+        head_indices = indices[idx_ptr[:-1]]
+        # generate down code, order, inverse
+        code = code[:, head_indices]
+        order = torch.argsort(code)
+        inverse = torch.zeros_like(order).scatter_(
+            dim=1,
+            index=order,
+            src=torch.arange(0, code.shape[1], device=order.device).repeat(
+                code.shape[0], 1
+            ),
+        )
+
+        if self.shuffle_orders:
+            perm = torch.randperm(code.shape[0])
+            code = code[perm]
+            order = order[perm]
+            inverse = inverse[perm]
+
+        # collect information
+        point_dict = Dict(
+            feat=torch_scatter.segment_csr(
+                self.proj(point.feat)[indices], idx_ptr, reduce=self.reduce
+            ),
+            coord=torch_scatter.segment_csr(
+                point.coord[indices], idx_ptr, reduce="mean"
+            ),
+            grid_coord=point.grid_coord[head_indices] >> pooling_depth,
+            serialized_code=code,
+            serialized_order=order,
+            serialized_inverse=inverse,
+            serialized_depth=point.serialized_depth - pooling_depth,
+            batch=point.batch[head_indices],
+        )
+
+        if "condition" in point.keys():
+            point_dict["condition"] = point.condition
+        if "context" in point.keys():
+            point_dict["context"] = point.context
+
+        if self.traceable:
+            point_dict["pooling_inverse"] = cluster
+            point_dict["pooling_parent"] = point
+        point = Point(point_dict)
+        if self.norm is not None:
+            point = self.norm(point)
+        if self.act is not None:
+            point = self.act(point)
+        point.sparsify()
+        return point
+
+
+class SerializedUnpooling(PointModule):
+    def __init__(
+        self,
+        in_channels,
+        skip_channels,
+        out_channels,
+        norm_layer=None,
+        act_layer=None,
+        traceable=False,  # record parent and cluster
+    ):
+        super().__init__()
+        self.proj = PointSequential(nn.Linear(in_channels, out_channels))
+        self.proj_skip = PointSequential(nn.Linear(skip_channels, out_channels))
+
+        if norm_layer is not None:
+            self.proj.add(norm_layer(out_channels))
+            self.proj_skip.add(norm_layer(out_channels))
+
+        if act_layer is not None:
+            self.proj.add(act_layer())
+            self.proj_skip.add(act_layer())
+
+        self.traceable = traceable
+
+    def forward(self, point):
+        assert "pooling_parent" in point.keys()
+        assert "pooling_inverse" in point.keys()
+        parent = point.pop("pooling_parent")
+        inverse = point.pop("pooling_inverse")
+        point = self.proj(point)
+        parent = self.proj_skip(parent)
+        parent.feat = parent.feat + point.feat[inverse]
+
+        if self.traceable:
+            parent["unpooling_parent"] = point
+        return parent
+
+
+class Embedding(PointModule):
+    def __init__(
+        self,
+        in_channels,
+        embed_channels,
+        norm_layer=None,
+        act_layer=None,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.embed_channels = embed_channels
+
+        # TODO: check remove spconv
+        self.stem = PointSequential(
+            conv=spconv.SubMConv3d(
+                in_channels,
+                embed_channels,
+                kernel_size=5,
+                padding=1,
+                bias=False,
+                indice_key="stem",
+            )
+        )
+        if norm_layer is not None:
+            self.stem.add(norm_layer(embed_channels), name="norm")
+        if act_layer is not None:
+            self.stem.add(act_layer(), name="act")
+
+    def forward(self, point: Point):
+        point = self.stem(point)
+        return point
+
+
+class PointTransformerV3(PointModule):
+    def __init__(
+        self,
+        in_channels=6,
+        order=("z", "z-trans", "hilbert", "hilbert-trans"),
+        stride=(2, 2, 2, 2),
+        enc_depths=(2, 2, 2, 6, 2),
+        enc_channels=(32, 64, 128, 256, 512),
+        enc_num_head=(2, 4, 8, 16, 32),
+        enc_patch_size=(1024, 1024, 1024, 1024, 1024),
+        dec_depths=(2, 2, 2, 2),
+        dec_channels=(64, 64, 128, 256),
+        dec_num_head=(4, 4, 8, 16),
+        dec_patch_size=(1024, 1024, 1024, 1024),
+        mlp_ratio=4,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        drop_path=0.3,
+        pre_norm=True,
+        shuffle_orders=True,
+        enable_rpe=False,
+        enable_flash=True,
+        upcast_attention=False,
+        upcast_softmax=False,
+        cls_mode=False,
+        pdnorm_bn=False,
+        pdnorm_ln=False,
+        pdnorm_decouple=True,
+        pdnorm_adaptive=False,
+        pdnorm_affine=True,
+        pdnorm_conditions=("ScanNet", "S3DIS", "Structured3D"),
+    ):
+        super().__init__()
+        self.num_stages = len(enc_depths)
+        self.order = [order] if isinstance(order, str) else order
+        self.cls_mode = cls_mode
+        self.shuffle_orders = shuffle_orders
+
+        assert self.num_stages == len(stride) + 1
+        assert self.num_stages == len(enc_depths)
+        assert self.num_stages == len(enc_channels)
+        assert self.num_stages == len(enc_num_head)
+        assert self.num_stages == len(enc_patch_size)
+        assert self.cls_mode or self.num_stages == len(dec_depths) + 1
+        assert self.cls_mode or self.num_stages == len(dec_channels) + 1
+        assert self.cls_mode or self.num_stages == len(dec_num_head) + 1
+        assert self.cls_mode or self.num_stages == len(dec_patch_size) + 1
+
+        # norm layers
+        if pdnorm_bn:
+            bn_layer = partial(
+                PDNorm,
+                norm_layer=partial(
+                    nn.BatchNorm1d, eps=1e-3, momentum=0.01, affine=pdnorm_affine
+                ),
+                conditions=pdnorm_conditions,
+                decouple=pdnorm_decouple,
+                adaptive=pdnorm_adaptive,
+            )
+        else:
+            bn_layer = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
+        if pdnorm_ln:
+            ln_layer = partial(
+                PDNorm,
+                norm_layer=partial(nn.LayerNorm, elementwise_affine=pdnorm_affine),
+                conditions=pdnorm_conditions,
+                decouple=pdnorm_decouple,
+                adaptive=pdnorm_adaptive,
+            )
+        else:
+            ln_layer = nn.LayerNorm
+        # activation layers
+        act_layer = nn.GELU
+
+        self.embedding = Embedding(
+            in_channels=in_channels,
+            embed_channels=enc_channels[0],
+            norm_layer=bn_layer,
+            act_layer=act_layer,
+        )
+
+        # encoder
+        enc_drop_path = [
+            x.item() for x in torch.linspace(0, drop_path, sum(enc_depths))
+        ]
+        self.enc = PointSequential()
+        for s in range(self.num_stages):
+            enc_drop_path_ = enc_drop_path[
+                sum(enc_depths[:s]) : sum(enc_depths[: s + 1])
+            ]
+            enc = PointSequential()
+            if s > 0:
+                enc.add(
+                    SerializedPooling(
+                        in_channels=enc_channels[s - 1],
+                        out_channels=enc_channels[s],
+                        stride=stride[s - 1],
+                        norm_layer=bn_layer,
+                        act_layer=act_layer,
+                    ),
+                    name="down",
+                )
+            for i in range(enc_depths[s]):
+                enc.add(
+                    Block(
+                        channels=enc_channels[s],
+                        num_heads=enc_num_head[s],
+                        patch_size=enc_patch_size[s],
+                        mlp_ratio=mlp_ratio,
+                        qkv_bias=qkv_bias,
+                        qk_scale=qk_scale,
+                        attn_drop=attn_drop,
+                        proj_drop=proj_drop,
+                        drop_path=enc_drop_path_[i],
+                        norm_layer=ln_layer,
+                        act_layer=act_layer,
+                        pre_norm=pre_norm,
+                        order_index=i % len(self.order),
+                        cpe_indice_key=f"stage{s}",
+                        enable_rpe=enable_rpe,
+                        enable_flash=enable_flash,
+                        upcast_attention=upcast_attention,
+                        upcast_softmax=upcast_softmax,
+                    ),
+                    name=f"block{i}",
+                )
+            if len(enc) != 0:
+                self.enc.add(module=enc, name=f"enc{s}")
+
+        # decoder
+        if not self.cls_mode:
+            dec_drop_path = [
+                x.item() for x in torch.linspace(0, drop_path, sum(dec_depths))
+            ]
+            self.dec = PointSequential()
+            dec_channels = list(dec_channels) + [enc_channels[-1]]
+            for s in reversed(range(self.num_stages - 1)):
+                dec_drop_path_ = dec_drop_path[
+                    sum(dec_depths[:s]) : sum(dec_depths[: s + 1])
+                ]
+                dec_drop_path_.reverse()
+                dec = PointSequential()
+                dec.add(
+                    SerializedUnpooling(
+                        in_channels=dec_channels[s + 1],
+                        skip_channels=enc_channels[s],
+                        out_channels=dec_channels[s],
+                        norm_layer=bn_layer,
+                        act_layer=act_layer,
+                    ),
+                    name="up",
+                )
+                for i in range(dec_depths[s]):
+                    dec.add(
+                        Block(
+                            channels=dec_channels[s],
+                            num_heads=dec_num_head[s],
+                            patch_size=dec_patch_size[s],
+                            mlp_ratio=mlp_ratio,
+                            qkv_bias=qkv_bias,
+                            qk_scale=qk_scale,
+                            attn_drop=attn_drop,
+                            proj_drop=proj_drop,
+                            drop_path=dec_drop_path_[i],
+                            norm_layer=ln_layer,
+                            act_layer=act_layer,
+                            pre_norm=pre_norm,
+                            order_index=i % len(self.order),
+                            cpe_indice_key=f"stage{s}",
+                            enable_rpe=enable_rpe,
+                            enable_flash=enable_flash,
+                            upcast_attention=upcast_attention,
+                            upcast_softmax=upcast_softmax,
+                        ),
+                        name=f"block{i}",
+                    )
+                self.dec.add(module=dec, name=f"dec{s}")
+
+    def forward(self, data_dict):
+        """
+        A data_dict is a dictionary containing properties of a batched point cloud.
+        It should contain the following properties for PTv3:
+        1. "feat": feature of point cloud
+        2. "grid_coord": discrete coordinate after grid sampling (voxelization) or "coord" + "grid_size"
+        3. "offset" or "batch": https://github.com/Pointcept/Pointcept?tab=readme-ov-file#offset
+        """
+        point = Point(data_dict)
+        point.serialization(order=self.order, shuffle_orders=self.shuffle_orders)
+        point.sparsify()
+
+        point = self.embedding(point)
+        point = self.enc(point)
+        if not self.cls_mode:
+            point = self.dec(point)
+        return point

experiments/ptv3/ptv3_code/serialization/__init__.py

@@ -0,0 +1,8 @@
+from .default import (
+    encode,
+    decode,
+    z_order_encode,
+    z_order_decode,
+    hilbert_encode,
+    hilbert_decode,
+)

experiments/ptv3/ptv3_code/serialization/default.py

@@ -0,0 +1,59 @@
+import torch
+from .z_order import xyz2key as z_order_encode_
+from .z_order import key2xyz as z_order_decode_
+from .hilbert import encode as hilbert_encode_
+from .hilbert import decode as hilbert_decode_
+
+
+@torch.inference_mode()
+def encode(grid_coord, batch=None, depth=16, order="z"):
+    assert order in {"z", "z-trans", "hilbert", "hilbert-trans"}
+    if order == "z":
+        code = z_order_encode(grid_coord, depth=depth)
+    elif order == "z-trans":
+        code = z_order_encode(grid_coord[:, [1, 0, 2]], depth=depth)
+    elif order == "hilbert":
+        code = hilbert_encode(grid_coord, depth=depth)
+    elif order == "hilbert-trans":
+        code = hilbert_encode(grid_coord[:, [1, 0, 2]], depth=depth)
+    else:
+        raise NotImplementedError
+    if batch is not None:
+        batch = batch.long()
+        code = batch << depth * 3 | code
+    return code
+
+
+@torch.inference_mode()
+def decode(code, depth=16, order="z"):
+    assert order in {"z", "hilbert"}
+    batch = code >> depth * 3
+    code = code & ((1 << depth * 3) - 1)
+    if order == "z":
+        grid_coord = z_order_decode(code, depth=depth)
+    elif order == "hilbert":
+        grid_coord = hilbert_decode(code, depth=depth)
+    else:
+        raise NotImplementedError
+    return grid_coord, batch
+
+
+def z_order_encode(grid_coord: torch.Tensor, depth: int = 16):
+    x, y, z = grid_coord[:, 0].long(), grid_coord[:, 1].long(), grid_coord[:, 2].long()
+    # we block the support to batch, maintain batched code in Point class
+    code = z_order_encode_(x, y, z, b=None, depth=depth)
+    return code
+
+
+def z_order_decode(code: torch.Tensor, depth):
+    x, y, z = z_order_decode_(code, depth=depth)
+    grid_coord = torch.stack([x, y, z], dim=-1)  # (N,  3)
+    return grid_coord
+
+
+def hilbert_encode(grid_coord: torch.Tensor, depth: int = 16):
+    return hilbert_encode_(grid_coord, num_dims=3, num_bits=depth)
+
+
+def hilbert_decode(code: torch.Tensor, depth: int = 16):
+    return hilbert_decode_(code, num_dims=3, num_bits=depth)

experiments/ptv3/ptv3_code/serialization/hilbert.py

@@ -0,0 +1,303 @@
+"""
+Hilbert Order
+Modified from https://github.com/PrincetonLIPS/numpy-hilbert-curve
+
+Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com), Kaixin Xu
+Please cite our work if the code is helpful to you.
+"""
+
+import torch
+
+
+def right_shift(binary, k=1, axis=-1):
+    """Right shift an array of binary values.
+
+    Parameters:
+    -----------
+     binary: An ndarray of binary values.
+
+     k: The number of bits to shift. Default 1.
+
+     axis: The axis along which to shift.  Default -1.
+
+    Returns:
+    --------
+     Returns an ndarray with zero prepended and the ends truncated, along
+     whatever axis was specified."""
+
+    # If we're shifting the whole thing, just return zeros.
+    if binary.shape[axis] <= k:
+        return torch.zeros_like(binary)
+
+    # Determine the padding pattern.
+    # padding = [(0,0)] * len(binary.shape)
+    # padding[axis] = (k,0)
+
+    # Determine the slicing pattern to eliminate just the last one.
+    slicing = [slice(None)] * len(binary.shape)
+    slicing[axis] = slice(None, -k)
+    shifted = torch.nn.functional.pad(
+        binary[tuple(slicing)], (k, 0), mode="constant", value=0
+    )
+
+    return shifted
+
+
+def binary2gray(binary, axis=-1):
+    """Convert an array of binary values into Gray codes.
+
+    This uses the classic X ^ (X >> 1) trick to compute the Gray code.
+
+    Parameters:
+    -----------
+     binary: An ndarray of binary values.
+
+     axis: The axis along which to compute the gray code. Default=-1.
+
+    Returns:
+    --------
+     Returns an ndarray of Gray codes.
+    """
+    shifted = right_shift(binary, axis=axis)
+
+    # Do the X ^ (X >> 1) trick.
+    gray = torch.logical_xor(binary, shifted)
+
+    return gray
+
+
+def gray2binary(gray, axis=-1):
+    """Convert an array of Gray codes back into binary values.
+
+    Parameters:
+    -----------
+     gray: An ndarray of gray codes.
+
+     axis: The axis along which to perform Gray decoding. Default=-1.
+
+    Returns:
+    --------
+     Returns an ndarray of binary values.
+    """
+
+    # Loop the log2(bits) number of times necessary, with shift and xor.
+    shift = 2 ** (torch.Tensor([gray.shape[axis]]).log2().ceil().int() - 1)
+    while shift > 0:
+        gray = torch.logical_xor(gray, right_shift(gray, shift))
+        shift = torch.div(shift, 2, rounding_mode="floor")
+    return gray
+
+
+def encode(locs, num_dims, num_bits):
+    """Decode an array of locations in a hypercube into a Hilbert integer.
+
+    This is a vectorized-ish version of the Hilbert curve implementation by John
+    Skilling as described in:
+
+    Skilling, J. (2004, April). Programming the Hilbert curve. In AIP Conference
+      Proceedings (Vol. 707, No. 1, pp. 381-387). American Institute of Physics.
+
+    Params:
+    -------
+     locs - An ndarray of locations in a hypercube of num_dims dimensions, in
+            which each dimension runs from 0 to 2**num_bits-1.  The shape can
+            be arbitrary, as long as the last dimension of the same has size
+            num_dims.
+
+     num_dims - The dimensionality of the hypercube. Integer.
+
+     num_bits - The number of bits for each dimension. Integer.
+
+    Returns:
+    --------
+     The output is an ndarray of uint64 integers with the same shape as the
+     input, excluding the last dimension, which needs to be num_dims.
+    """
+
+    # Keep around the original shape for later.
+    orig_shape = locs.shape
+    bitpack_mask = 1 << torch.arange(0, 8).to(locs.device)
+    bitpack_mask_rev = bitpack_mask.flip(-1)
+
+    if orig_shape[-1] != num_dims:
+        raise ValueError(
+            """
+      The shape of locs was surprising in that the last dimension was of size
+      %d, but num_dims=%d.  These need to be equal.
+      """
+            % (orig_shape[-1], num_dims)
+        )
+
+    if num_dims * num_bits > 63:
+        raise ValueError(
+            """
+      num_dims=%d and num_bits=%d for %d bits total, which can't be encoded
+      into a int64.  Are you sure you need that many points on your Hilbert
+      curve?
+      """
+            % (num_dims, num_bits, num_dims * num_bits)
+        )
+
+    # Treat the location integers as 64-bit unsigned and then split them up into
+    # a sequence of uint8s.  Preserve the association by dimension.
+    locs_uint8 = locs.long().view(torch.uint8).reshape((-1, num_dims, 8)).flip(-1)
+
+    # Now turn these into bits and truncate to num_bits.
+    gray = (
+        locs_uint8.unsqueeze(-1)
+        .bitwise_and(bitpack_mask_rev)
+        .ne(0)
+        .byte()
+        .flatten(-2, -1)[..., -num_bits:]
+    )
+
+    # Run the decoding process the other way.
+    # Iterate forwards through the bits.
+    for bit in range(0, num_bits):
+        # Iterate forwards through the dimensions.
+        for dim in range(0, num_dims):
+            # Identify which ones have this bit active.
+            mask = gray[:, dim, bit]
+
+            # Where this bit is on, invert the 0 dimension for lower bits.
+            gray[:, 0, bit + 1 :] = torch.logical_xor(
+                gray[:, 0, bit + 1 :], mask[:, None]
+            )
+
+            # Where the bit is off, exchange the lower bits with the 0 dimension.
+            to_flip = torch.logical_and(
+                torch.logical_not(mask[:, None]).repeat(1, gray.shape[2] - bit - 1),
+                torch.logical_xor(gray[:, 0, bit + 1 :], gray[:, dim, bit + 1 :]),
+            )
+            gray[:, dim, bit + 1 :] = torch.logical_xor(
+                gray[:, dim, bit + 1 :], to_flip
+            )
+            gray[:, 0, bit + 1 :] = torch.logical_xor(gray[:, 0, bit + 1 :], to_flip)
+
+    # Now flatten out.
+    gray = gray.swapaxes(1, 2).reshape((-1, num_bits * num_dims))
+
+    # Convert Gray back to binary.
+    hh_bin = gray2binary(gray)
+
+    # Pad back out to 64 bits.
+    extra_dims = 64 - num_bits * num_dims
+    padded = torch.nn.functional.pad(hh_bin, (extra_dims, 0), "constant", 0)
+
+    # Convert binary values into uint8s.
+    hh_uint8 = (
+        (padded.flip(-1).reshape((-1, 8, 8)) * bitpack_mask)
+        .sum(2)
+        .squeeze()
+        .type(torch.uint8)
+    )
+
+    # Convert uint8s into uint64s.
+    hh_uint64 = hh_uint8.view(torch.int64).squeeze()
+
+    return hh_uint64
+
+
+def decode(hilberts, num_dims, num_bits):
+    """Decode an array of Hilbert integers into locations in a hypercube.
+
+    This is a vectorized-ish version of the Hilbert curve implementation by John
+    Skilling as described in:
+
+    Skilling, J. (2004, April). Programming the Hilbert curve. In AIP Conference
+      Proceedings (Vol. 707, No. 1, pp. 381-387). American Institute of Physics.
+
+    Params:
+    -------
+     hilberts - An ndarray of Hilbert integers.  Must be an integer dtype and
+                cannot have fewer bits than num_dims * num_bits.
+
+     num_dims - The dimensionality of the hypercube. Integer.
+
+     num_bits - The number of bits for each dimension. Integer.
+
+    Returns:
+    --------
+     The output is an ndarray of unsigned integers with the same shape as hilberts
+     but with an additional dimension of size num_dims.
+    """
+
+    if num_dims * num_bits > 64:
+        raise ValueError(
+            """
+      num_dims=%d and num_bits=%d for %d bits total, which can't be encoded
+      into a uint64.  Are you sure you need that many points on your Hilbert
+      curve?
+      """
+            % (num_dims, num_bits)
+        )
+
+    # Handle the case where we got handed a naked integer.
+    hilberts = torch.atleast_1d(hilberts)
+
+    # Keep around the shape for later.
+    orig_shape = hilberts.shape
+    bitpack_mask = 2 ** torch.arange(0, 8).to(hilberts.device)
+    bitpack_mask_rev = bitpack_mask.flip(-1)
+
+    # Treat each of the hilberts as a s equence of eight uint8.
+    # This treats all of the inputs as uint64 and makes things uniform.
+    hh_uint8 = (
+        hilberts.ravel().type(torch.int64).view(torch.uint8).reshape((-1, 8)).flip(-1)
+    )
+
+    # Turn these lists of uints into lists of bits and then truncate to the size
+    # we actually need for using Skilling's procedure.
+    hh_bits = (
+        hh_uint8.unsqueeze(-1)
+        .bitwise_and(bitpack_mask_rev)
+        .ne(0)
+        .byte()
+        .flatten(-2, -1)[:, -num_dims * num_bits :]
+    )
+
+    # Take the sequence of bits and Gray-code it.
+    gray = binary2gray(hh_bits)
+
+    # There has got to be a better way to do this.
+    # I could index them differently, but the eventual packbits likes it this way.
+    gray = gray.reshape((-1, num_bits, num_dims)).swapaxes(1, 2)
+
+    # Iterate backwards through the bits.
+    for bit in range(num_bits - 1, -1, -1):
+        # Iterate backwards through the dimensions.
+        for dim in range(num_dims - 1, -1, -1):
+            # Identify which ones have this bit active.
+            mask = gray[:, dim, bit]
+
+            # Where this bit is on, invert the 0 dimension for lower bits.
+            gray[:, 0, bit + 1 :] = torch.logical_xor(
+                gray[:, 0, bit + 1 :], mask[:, None]
+            )
+
+            # Where the bit is off, exchange the lower bits with the 0 dimension.
+            to_flip = torch.logical_and(
+                torch.logical_not(mask[:, None]),
+                torch.logical_xor(gray[:, 0, bit + 1 :], gray[:, dim, bit + 1 :]),
+            )
+            gray[:, dim, bit + 1 :] = torch.logical_xor(
+                gray[:, dim, bit + 1 :], to_flip
+            )
+            gray[:, 0, bit + 1 :] = torch.logical_xor(gray[:, 0, bit + 1 :], to_flip)
+
+    # Pad back out to 64 bits.
+    extra_dims = 64 - num_bits
+    padded = torch.nn.functional.pad(gray, (extra_dims, 0), "constant", 0)
+
+    # Now chop these up into blocks of 8.
+    locs_chopped = padded.flip(-1).reshape((-1, num_dims, 8, 8))
+
+    # Take those blocks and turn them unto uint8s.
+    # from IPython import embed; embed()
+    locs_uint8 = (locs_chopped * bitpack_mask).sum(3).squeeze().type(torch.uint8)
+
+    # Finally, treat these as uint64s.
+    flat_locs = locs_uint8.view(torch.int64)
+
+    # Return them in the expected shape.
+    return flat_locs.reshape((*orig_shape, num_dims))

experiments/ptv3/ptv3_code/serialization/z_order.py

@@ -0,0 +1,126 @@
+# --------------------------------------------------------
+# Octree-based Sparse Convolutional Neural Networks
+# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Peng-Shuai Wang
+# --------------------------------------------------------
+
+import torch
+from typing import Optional, Union
+
+
+class KeyLUT:
+    def __init__(self):
+        r256 = torch.arange(256, dtype=torch.int64)
+        r512 = torch.arange(512, dtype=torch.int64)
+        zero = torch.zeros(256, dtype=torch.int64)
+        device = torch.device("cpu")
+
+        self._encode = {
+            device: (
+                self.xyz2key(r256, zero, zero, 8),
+                self.xyz2key(zero, r256, zero, 8),
+                self.xyz2key(zero, zero, r256, 8),
+            )
+        }
+        self._decode = {device: self.key2xyz(r512, 9)}
+
+    def encode_lut(self, device=torch.device("cpu")):
+        if device not in self._encode:
+            cpu = torch.device("cpu")
+            self._encode[device] = tuple(e.to(device) for e in self._encode[cpu])
+        return self._encode[device]
+
+    def decode_lut(self, device=torch.device("cpu")):
+        if device not in self._decode:
+            cpu = torch.device("cpu")
+            self._decode[device] = tuple(e.to(device) for e in self._decode[cpu])
+        return self._decode[device]
+
+    def xyz2key(self, x, y, z, depth):
+        key = torch.zeros_like(x)
+        for i in range(depth):
+            mask = 1 << i
+            key = (
+                key
+                | ((x & mask) << (2 * i + 2))
+                | ((y & mask) << (2 * i + 1))
+                | ((z & mask) << (2 * i + 0))
+            )
+        return key
+
+    def key2xyz(self, key, depth):
+        x = torch.zeros_like(key)
+        y = torch.zeros_like(key)
+        z = torch.zeros_like(key)
+        for i in range(depth):
+            x = x | ((key & (1 << (3 * i + 2))) >> (2 * i + 2))
+            y = y | ((key & (1 << (3 * i + 1))) >> (2 * i + 1))
+            z = z | ((key & (1 << (3 * i + 0))) >> (2 * i + 0))
+        return x, y, z
+
+
+_key_lut = KeyLUT()
+
+
+def xyz2key(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    z: torch.Tensor,
+    b: Optional[Union[torch.Tensor, int]] = None,
+    depth: int = 16,
+):
+    r"""Encodes :attr:`x`, :attr:`y`, :attr:`z` coordinates to the shuffled keys
+    based on pre-computed look up tables. The speed of this function is much
+    faster than the method based on for-loop.
+
+    Args:
+      x (torch.Tensor): The x coordinate.
+      y (torch.Tensor): The y coordinate.
+      z (torch.Tensor): The z coordinate.
+      b (torch.Tensor or int): The batch index of the coordinates, and should be
+          smaller than 32768. If :attr:`b` is :obj:`torch.Tensor`, the size of
+          :attr:`b` must be the same as :attr:`x`, :attr:`y`, and :attr:`z`.
+      depth (int): The depth of the shuffled key, and must be smaller than 17 (< 17).
+    """
+
+    EX, EY, EZ = _key_lut.encode_lut(x.device)
+    x, y, z = x.long(), y.long(), z.long()
+
+    mask = 255 if depth > 8 else (1 << depth) - 1
+    key = EX[x & mask] | EY[y & mask] | EZ[z & mask]
+    if depth > 8:
+        mask = (1 << (depth - 8)) - 1
+        key16 = EX[(x >> 8) & mask] | EY[(y >> 8) & mask] | EZ[(z >> 8) & mask]
+        key = key16 << 24 | key
+
+    if b is not None:
+        b = b.long()
+        key = b << 48 | key
+
+    return key
+
+
+def key2xyz(key: torch.Tensor, depth: int = 16):
+    r"""Decodes the shuffled key to :attr:`x`, :attr:`y`, :attr:`z` coordinates
+    and the batch index based on pre-computed look up tables.
+
+    Args:
+      key (torch.Tensor): The shuffled key.
+      depth (int): The depth of the shuffled key, and must be smaller than 17 (< 17).
+    """
+
+    DX, DY, DZ = _key_lut.decode_lut(key.device)
+    x, y, z = torch.zeros_like(key), torch.zeros_like(key), torch.zeros_like(key)
+
+    b = key >> 48
+    key = key & ((1 << 48) - 1)
+
+    n = (depth + 2) // 3
+    for i in range(n):
+        k = key >> (i * 9) & 511
+        x = x | (DX[k] << (i * 3))
+        y = y | (DY[k] << (i * 3))
+        z = z | (DZ[k] << (i * 3))
+
+    return x, y, z, b

experiments/ptv3/train_args.json

@@ -0,0 +1,62 @@
+{
+  "activation": "gelu",
+  "adam_betas": "0.9,0.95",
+  "arch": "ptv3",
+  "aug_drop": 0.0,
+  "aug_flip": true,
+  "aug_jitter": 0.0,
+  "aug_rotate": true,
+  "batch_size": 32,
+  "behind_emb_dim": 8,
+  "cache_dir": "hf://usm3d/s23dr-2026-sampled_8192_v3:train",
+  "conf_clamp_min": null,
+  "conf_head_wd": 0.1,
+  "conf_mode": "sinkhorn",
+  "conf_weight": 0.1,
+  "cooldown_start": 195000,
+  "cooldown_steps": 5000,
+  "cosine_decay": false,
+  "cpu": false,
+  "cross_attn_interval": 4,
+  "decoder_input_xattn": false,
+  "decoder_layers": 3,
+  "deterministic": false,
+  "dropout": 0.1,
+  "ema_decay": 0.0,
+  "encoder_layers": 4,
+  "endpoint_warmup": 0,
+  "endpoint_weight": 0.1,
+  "ff": 1024,
+  "hidden": 256,
+  "kv_heads_cross": 2,
+  "kv_heads_self": 2,
+  "latent_layers": 7,
+  "latent_tokens": 256,
+  "learnable_fourier": false,
+  "length_floor": 0.0,
+  "lr": 0.00015,
+  "num_heads": 4,
+  "out_dir": "runs/step3b_8192",
+  "pre_encoder_layers": 0,
+  "qk_norm": true,
+  "qk_norm_type": "l2",
+  "rms_norm": true,
+  "seed": 353,
+  "segment_conf": true,
+  "segment_param": "midpoint_dir_len",
+  "segments": 64,
+  "seq_len": 8192,
+  "sinkhorn_dustbin": 0.3,
+  "sinkhorn_eps": 0.1,
+  "sinkhorn_eps_schedule": "none",
+  "sinkhorn_eps_start": null,
+  "sinkhorn_iters": 20,
+  "sinkhorn_weight": 1.0,
+  "steps": 200000,
+  "val_cache_dir": "",
+  "varifold_cross_only": false,
+  "varifold_weight": 0.0,
+  "vote_features": true,
+  "warmup": 1000,
+  "weight_decay": 0.01
+}

pnet_class_2026.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0bab151af7133e4da4e00734c12145d7dcb882297fa5172a6b7ba737dc130b6
+size 22456481

requirements.txt

@@ -0,0 +1,12 @@
+numpy
+torch
+tqdm
+datasets
+huggingface_hub
+opencv-python
+scipy
+scikit-learn
+pillow
+pycolmap>0.6
+trimesh
+hoho2025

s23dr_2026_example/attention.py

@@ -0,0 +1,141 @@
+# custom_transformer.py
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# =============================================================================
+# Core Efficient Multihead Attention using Scaled Dot Product Attention (SDPA)
+# =============================================================================
+
+class MultiHeadSDPA(nn.Module):
+    """
+    Multi-head cross-attention using torch.nn.functional.scaled_dot_product_attention
+    without causal masking. Suitable for set inputs and cross-attention.
+
+    If qk_norm=True, L2-normalizes Q and K per-head before the dot product,
+    then scales by a learned per-head temperature (log_scale). This caps logit
+    magnitude to [-1, +1] * exp(log_scale), preventing attention entropy
+    collapse at large head_dim.
+    """
+    def __init__(self, d_model: int, num_heads: int, kv_heads: int = None,
+                 qk_norm: bool = False, qk_norm_type: str = "l2"):
+        super().__init__()
+        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.kv_heads = kv_heads or num_heads
+        assert self.num_heads % self.kv_heads == 0, "kv_heads must divide num_heads"
+
+        self.head_dim = d_model // num_heads
+        self.qk_norm = qk_norm
+        self.qk_norm_type = qk_norm_type
+
+        # Input projection layers
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, self.kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(d_model, self.kv_heads * self.head_dim, bias=False)
+
+        # Output projection
+        self.out_proj = nn.Linear(d_model, d_model, bias=False)
+        nn.init.zeros_(self.out_proj.weight)
+
+        if qk_norm:
+            import math
+            if qk_norm_type == "rms":
+                # Standard QK-norm (Qwen3/Gemma3 style): RMSNorm on Q and K,
+                # no learned temperature. SDPA's 1/sqrt(d) scaling is sufficient
+                # because RMSNorm preserves the expected logit variance.
+                pass  # no extra parameters needed
+            else:
+                # L2 + learned temperature (nGPT/ViT-22B style):
+                # L2 projects to unit sphere, needs learned scale to compensate.
+                self.log_scale = nn.Parameter(
+                    torch.full((num_heads,), math.log(math.sqrt(self.head_dim))))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # Project
+        q = self.q_proj(query)
+        k = self.k_proj(key)
+        v = self.v_proj(key)
+
+        B, Tq, _ = q.shape
+        _, Tk, _ = k.shape
+
+        q = q.view(B, Tq, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, Tk, self.kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, Tk, self.kv_heads, self.head_dim).transpose(1, 2)
+
+        if self.kv_heads != self.num_heads:
+            repeat = self.num_heads // self.kv_heads
+            k = k.repeat_interleave(repeat, dim=1)
+            v = v.repeat_interleave(repeat, dim=1)
+
+        if self.qk_norm:
+            if self.qk_norm_type == "rms":
+                # RMSNorm (Qwen3/Gemma3 style): no learned temperature needed.
+                # After RMSNorm, logit variance matches standard SDPA naturally.
+                q = q * torch.rsqrt(q.square().mean(dim=-1, keepdim=True) + 1e-6)
+                k = k * torch.rsqrt(k.square().mean(dim=-1, keepdim=True) + 1e-6)
+                attn_mask = None
+                if key_padding_mask is not None:
+                    attn_mask = ~key_padding_mask[:, None, None, :].to(dtype=torch.bool)
+                attn_out = F.scaled_dot_product_attention(
+                    q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False,
+                )
+            else:
+                # L2 + learned temperature (nGPT/ViT-22B style)
+                q = F.normalize(q, dim=-1)
+                k = F.normalize(k, dim=-1)
+                scale = self.log_scale.exp().view(1, -1, 1, 1)
+                q = q * scale
+                attn_mask = None
+                if key_padding_mask is not None:
+                    attn_mask = ~key_padding_mask[:, None, None, :].to(dtype=torch.bool)
+                attn_out = F.scaled_dot_product_attention(
+                    q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False,
+                    scale=1.0,
+                )
+        else:
+            attn_mask = None
+            if key_padding_mask is not None:
+                attn_mask = ~key_padding_mask[:, None, None, :].to(dtype=torch.bool)
+            attn_out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False
+            )
+
+        attn_out = attn_out.transpose(1, 2).reshape(B, Tq, self.d_model)
+        return self.out_proj(attn_out)
+
+
+# =============================================================================
+# Transformer Feed-Forward Block
+# =============================================================================
+
+def _get_activation(name: str):
+    """Look up activation function by name. Supports 'relu_sq' for ReLU^2."""
+    if name == "relu_sq":
+        return lambda x: F.relu(x).square()
+    return getattr(F, name)
+
+
+class FeedForward(nn.Module):
+    """
+    Position-wise MLP block: linear -> activation -> linear.
+    Supports 'gelu', 'relu', 'relu_sq', etc.
+    """
+    def __init__(self, d_model: int, dim_ff: int, activation: str = "gelu"):
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, dim_ff)
+        self.linear2 = nn.Linear(dim_ff, d_model)
+        nn.init.zeros_(self.linear2.weight)
+        nn.init.zeros_(self.linear2.bias)
+        self.activation = _get_activation(activation)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.linear1(x)
+        return self.linear2(self.activation(x))

s23dr_2026_example/bad_samples.txt

@@ -0,0 +1,156 @@
+14b1872e960
+1807ef90db4
+180e6a67e87
+1ad5c6bd31f
+1c3f939ad93
+1ede4c0d52f
+214f17d9cc4
+22256d88df9
+24a92a8de6d
+24b4e984bad
+2565978cf53
+2a71f1a2072
+2d44c1fade6
+2ebed43823a
+33982551420
+3b480496f82
+412a2bdf7a4
+44343bbabbb
+4a0b3f04cbd
+4a7fa170826
+4b7dc027214
+4e0dc2c9b18
+5172a516c8b
+529e8f15cd2
+56fc6f6f163
+575963ce814
+578ec40a278
+5a0c07c575a
+5d521223c26
+6148b5c9461
+631eb6d7c03
+655a14f8a75
+66502d7ee6f
+6da76fc6687
+777eaaad0ca
+7a4e2909d68
+7c5c9baf483
+80806dfd75e
+81a4ead431d
+833152dd554
+85797868c0f
+86460ad8181
+86783a6bee4
+95193322d7a
+99a9d056200
+9b1d4eeaab9
+9ff759f2e4c
+acbd243da16
+b9b275710c0
+beceaa9bb7c
+c243d079286
+c5c7337d2cb
+cdf6f2d3b35
+cfe370f1c87
+d4a72aea80c
+d655f066cd3
+d79e8d9455c
+d7d6c5be76e
+dc30ae4b93b
+de9495f7ca3
+e1901819c72
+e1d88c1a6b1
+e5d3eb0a617
+ec11d3cdcf6
+ecb21fad0ad
+ee55d8c6493
+ee7e6d4dee1
+008052054aa
+03ecb7d3cf3
+0555a655534
+099cad230c6
+0d061ae23f0
+10741a421c0
+110d5e407b9
+128a7fb415a
+13177736b26
+1635d73bf7d
+18a760de9ea
+18d90d03e95
+209627a5c1a
+21e3cd4b7b8
+22f5499200d
+266eb64de68
+269235f770b
+2758490e558
+2a203cf5d35
+2a878ec47ab
+2cb43eb2201
+393298e282b
+395abe6aac7
+3d19c7a4ca3
+44e2b719b1e
+45039819fcc
+4cb4ff01619
+4e5eb5712fa
+4e988765a6d
+5077bf42714
+55ed69b2622
+5ae3b651a37
+5ca1edeed4c
+5daa76b1c7f
+5fdd11dfae5
+6078cf180c2
+6682b309e9c
+6c02d2038c0
+71c595506c8
+73c8f960c18
+74ccc8fd057
+7a34156a798
+7ac7af9f59c
+7f2ec0ea179
+823b837b36c
+82d7600f9a3
+848161a2900
+88cedf129eb
+8dec106b6a6
+8e335d08ca4
+8ecf7c58193
+8fa55008beb
+90e09de2301
+9197acc0b9d
+954c25e876c
+98517d5563d
+99e717a0148
+9a0c0635bd7
+9ad436b7b3d
+9be351cbf14
+9e2a2e51798
+a84a7ea9220
+aa8cb84d3eb
+b07977292da
+b3e33456f0b
+b7823de373e
+bac379382d9
+bd2d9bf67a3
+c14584a84cd
+c497170c970
+cd8e767612b
+d17917bb279
+d42b9d432a9
+d53d8857a85
+d6808cf3d98
+d6f509d1dd9
+d7abd08e643
+d83493bf974
+d87293651ee
+da9d4ac9e8e
+daa1702791a
+dcb12411c14
+de9ab9cdd5b
+df906c58a3c
+e3870649eb5
+ea90aed9b98
+ecaa81b9711
+efc1238665b
+c5a65219daf

s23dr_2026_example/cache_scenes.py

@@ -0,0 +1,282 @@
+#!/usr/bin/env python3
+"""Cache compact scenes from HoHo22k shards to training-ready .pt files.
+
+Streams samples from the public `usm3d/hoho22k_2026_trainval` dataset, runs
+`build_compact_scene` (see point_fusion.py), precomputes priority group_id
+and semantic class_id, and saves one .pt per scene.
+
+Stage 1 of the dataset pipeline. See make_sampled_cache.py for stage 2.
+
+Usage:
+  python -m s23dr_2026_example.cache_scenes --out-dir cache/full --split train
+  python -m s23dr_2026_example.cache_scenes --out-dir cache/full_val --split validation
+
+Cache format per .pt file:
+  xyz:             float32 [P, 3]   all points in world space
+  source:          uint8   [P]      0=colmap, 1=depth
+  group_id:        int8    [P]      priority tier 0-4, -1=excluded
+  class_id:        uint8   [P]      one-hot class index (0-12)
+  behind_gest_id:  int16   [P]      behind-gestalt id (-1 if none)
+  visible_src:     uint8   [P]      1=gestalt, 2=ade
+  visible_id:      int16   [P]      class id within space
+  n_views_voted:   uint8   [P]      number of views that voted
+  vote_frac:       float32 [P]      fraction of votes
+  center:          float32 [3]      smart normalization center
+  scale:           float32 scalar   smart normalization scale
+  gt_vertices:     float32 [V, 3]   ground truth wireframe vertices
+  gt_edges:        int32   [E, 2]   ground truth wireframe edge indices
+"""
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .point_fusion import (
+    FuserConfig, build_compact_scene,
+    GEST_ID_TO_NAME, ADE_ID_TO_NAME, NUM_GEST,
+)
+
+# ---------------------------------------------------------------------------
+# Semantic class encoding: 11 structural + 1 other_house + 1 non_house = 13
+# ---------------------------------------------------------------------------
+
+# Each structural gestalt class gets its own one-hot bit.
+STRUCTURAL_CLASSES = (
+    "apex", "eave_end_point", "flashing_end_point",     # point classes (tier 0)
+    "rake", "ridge", "eave", "hip", "valley",           # roof edges (tier 1)
+    "flashing", "step_flashing",
+    "roof",                                              # roof face (tier 2)
+)
+# Index 11 = other house part (door, window, siding, etc.)
+# Index 12 = non-house / ADE / unlabeled
+NUM_SEMANTIC_CLASSES = len(STRUCTURAL_CLASSES) + 2  # 13
+
+# Priority tiers (same as tokenizer.py)
+_GEST_NAME_TO_ID = {n: i for i, n in enumerate(GEST_ID_TO_NAME)}
+_POINT_IDS = {_GEST_NAME_TO_ID[n] for n in ("apex", "eave_end_point", "flashing_end_point") if n in _GEST_NAME_TO_ID}
+_EDGE_IDS = {_GEST_NAME_TO_ID[n] for n in ("rake", "ridge", "eave", "hip", "valley", "flashing", "step_flashing") if n in _GEST_NAME_TO_ID}
+_FACE_IDS = {_GEST_NAME_TO_ID[n] for n in ("roof",) if n in _GEST_NAME_TO_ID}
+_HOUSE_IDS = {_GEST_NAME_TO_ID[n] for n in (
+    "apex", "eave_end_point", "flashing_end_point",
+    "rake", "ridge", "eave", "hip", "valley", "flashing", "step_flashing",
+    "roof", "door", "garage", "window", "shutter", "fascia", "soffit",
+    "horizontal_siding", "vertical_siding", "brick", "concrete",
+    "other_wall", "trim", "post", "ground_line",
+) if n in _GEST_NAME_TO_ID}
+
+_ADE_NAME_TO_ID = {n.lower(): i for i, n in enumerate(ADE_ID_TO_NAME)}
+_ADE_HOUSE_IDS = {_ADE_NAME_TO_ID[n] for n in ("building;edifice", "house", "wall", "windowpane;window", "door;double;door") if n in _ADE_NAME_TO_ID}
+
+_UNCLS_ID = _GEST_NAME_TO_ID.get("unclassified", -1)
+
+# Map structural gestalt names to one-hot index
+_STRUCTURAL_ONEHOT = {}
+for idx, name in enumerate(STRUCTURAL_CLASSES):
+    gid = _GEST_NAME_TO_ID.get(name)
+    if gid is not None:
+        _STRUCTURAL_ONEHOT[gid] = idx
+
+
+def _compute_group_and_class(visible_src, visible_id, behind_id, source):
+    """Compute priority group_id and semantic class_id per point (vectorized).
+
+    Args:
+        visible_src: uint8 [P] -- 0=unlabeled, 1=gestalt, 2=ade
+        visible_id:  int16 [P] -- class id within gestalt or ade space
+        behind_id:   int16 [P] -- behind-gestalt id (-1 if none)
+        source:      uint8 [P] -- 0=colmap, 1=depth
+
+    Returns:
+        group_id:  int8  [P] -- priority tier 0-4, -1 for excluded (unclassified)
+        class_id:  uint8 [P] -- one-hot class index 0-12
+    """
+    P = len(visible_src)
+    vsrc = visible_src.astype(np.int32)
+    vid = visible_id.astype(np.int32)
+    bid = behind_id.astype(np.int32)
+
+    # Effective gestalt id: prefer visible gestalt, fall back to behind
+    gest_id = np.full(P, -1, dtype=np.int32)
+    has_vis_gest = (vsrc == 1) & (vid >= 0)
+    has_behind = (bid >= 0) & ~has_vis_gest
+    gest_id[has_vis_gest] = vid[has_vis_gest]
+    gest_id[has_behind] = bid[has_behind]
+
+    # Exclude unclassified points
+    if _UNCLS_ID >= 0:
+        is_uncls = ((vsrc == 1) & (vid == _UNCLS_ID)) | (bid == _UNCLS_ID)
+        gest_id[is_uncls] = -1  # force excluded
+
+    # Build lookup arrays for gestalt id -> group and gestalt id -> class
+    max_gid = NUM_GEST
+    gid_to_group = np.full(max_gid, 4, dtype=np.int8)  # default: tier 4
+    gid_to_class = np.full(max_gid, NUM_SEMANTIC_CLASSES - 1, dtype=np.uint8)  # default: non-house
+
+    for gid in _POINT_IDS:
+        gid_to_group[gid] = 0
+    for gid in _EDGE_IDS:
+        gid_to_group[gid] = 1
+    for gid in _FACE_IDS:
+        gid_to_group[gid] = 2
+    for gid in _HOUSE_IDS - _POINT_IDS - _EDGE_IDS - _FACE_IDS:
+        gid_to_group[gid] = 3
+    for gid, onehot_idx in _STRUCTURAL_ONEHOT.items():
+        gid_to_class[gid] = onehot_idx
+    for gid in _HOUSE_IDS - set(_STRUCTURAL_ONEHOT.keys()):
+        gid_to_class[gid] = len(STRUCTURAL_CLASSES)  # other_house
+
+    # Apply lookup for points with valid gestalt ids
+    has_gest = gest_id >= 0
+    group_id = np.full(P, 4, dtype=np.int8)  # default: tier 4
+    class_id = np.full(P, NUM_SEMANTIC_CLASSES - 1, dtype=np.uint8)  # default: non-house
+
+    group_id[has_gest] = gid_to_group[gest_id[has_gest]]
+    class_id[has_gest] = gid_to_class[gest_id[has_gest]]
+
+    # ADE house points (no gestalt) get tier 3 + class_id = other_house
+    ade_house_arr = np.array(sorted(_ADE_HOUSE_IDS), dtype=np.int32)
+    is_ade_house = ~has_gest & (vsrc == 2) & (vid >= 0) & np.isin(vid, ade_house_arr)
+    group_id[is_ade_house] = 3
+    class_id[is_ade_house] = len(STRUCTURAL_CLASSES)  # other_house (index 11)
+
+    # Mark excluded points (unclassified) as -1
+    if _UNCLS_ID >= 0:
+        group_id[is_uncls] = -1
+        class_id[is_uncls] = NUM_SEMANTIC_CLASSES - 1
+
+    return group_id, class_id
+
+
+def _compute_smart_center_scale(xyz, source, mad_k=2.5, percentile=95.0,
+                                 max_points=8000):
+    """Compute normalization center and scale from depth points with MAD filter."""
+    depth_mask = source == 1
+    ref = xyz[depth_mask] if depth_mask.any() else xyz
+    if ref.shape[0] == 0:
+        center = xyz.mean(axis=0)
+        scale = max(np.linalg.norm(xyz - center, axis=1).max(), 1e-6)
+        return center.astype(np.float32), np.float32(scale)
+
+    if ref.shape[0] > max_points:
+        idx = np.random.choice(ref.shape[0], max_points, replace=False)
+        ref = ref[idx]
+
+    center0 = np.median(ref, axis=0)
+    dist = np.linalg.norm(ref - center0, axis=1)
+    med = np.median(dist)
+    mad = max(np.median(np.abs(dist - med)), 1e-6)
+    inliers = dist <= (med + mad_k * mad)
+    if inliers.any():
+        ref = ref[inliers]
+
+    # Percentile bounding box
+    lo_f = (100.0 - percentile) * 0.5 / 100.0
+    sorted_v = np.sort(ref, axis=0)
+    n = sorted_v.shape[0]
+    lo_idx = max(0, min(n - 1, int(lo_f * (n - 1))))
+    hi_idx = max(0, min(n - 1, int((1.0 - lo_f) * (n - 1))))
+    low = sorted_v[lo_idx]
+    high = sorted_v[hi_idx]
+
+    center = 0.5 * (low + high)
+    scale = max(np.sqrt(((high - low) ** 2).sum()), 1e-6)
+    return center.astype(np.float32), np.float32(scale)
+
+
+# ---------------------------------------------------------------------------
+# Dataset pipeline stage 1: raw HF sample -> cached .pt
+# ---------------------------------------------------------------------------
+
+def _process_one(sample, cfg):
+    """Fuse a single HF sample into a cache dict. Returns (order_id, dict) or None."""
+    rng = np.random.RandomState()
+
+    n_edges = len(sample.get("wf_edges", []))
+    if n_edges == 0 or n_edges > 64:
+        return None
+
+    scene = build_compact_scene(sample, cfg, rng=rng)
+    if scene is None:
+        return None
+
+    gt_v = scene.get("gt_vertices")
+    gt_e = scene.get("gt_edges")
+    if gt_v is None or gt_e is None or len(gt_e) == 0:
+        return None
+
+    xyz = scene["xyz"]
+    source = scene["source"]
+    group_id, class_id = _compute_group_and_class(
+        scene["visible_src"], scene["visible_id"], scene["behind_gest_id"], source)
+    center, scale = _compute_smart_center_scale(xyz, source)
+
+    gt_edge_classes = np.asarray(sample["wf_classifications"], dtype=np.int64)
+    return sample["order_id"], {
+        "xyz": xyz.astype(np.float32),
+        "source": source.astype(np.uint8),
+        "group_id": group_id,
+        "class_id": class_id,
+        "behind_gest_id": scene["behind_gest_id"].astype(np.int16),
+        "visible_src": scene["visible_src"].astype(np.uint8),
+        "visible_id": scene["visible_id"].astype(np.int16),
+        "n_views_voted": scene["n_views_voted"],
+        "vote_frac": scene["vote_frac"],
+        "center": center,
+        "scale": scale,
+        "gt_vertices": gt_v.astype(np.float32),
+        "gt_edges": gt_e.astype(np.int32),
+        "gt_edge_classes": gt_edge_classes,
+    }
+
+
+def main():
+    p = argparse.ArgumentParser(description="Stage 1: HoHo22k -> cached .pt files")
+    p.add_argument("--out-dir", required=True, help="Output directory for .pt files")
+    p.add_argument("--split", default="train", choices=["train", "validation"])
+    p.add_argument("--limit", type=int, default=0, help="Stop after N samples (0 = all)")
+    p.add_argument("--depth-per-view", type=int, default=8000)
+    p.add_argument("--skip-existing", action="store_true")
+    args = p.parse_args()
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    existing = {p.stem for p in out_dir.glob("*.pt")} if args.skip_existing else set()
+
+    from datasets import load_dataset
+    print(f"Streaming usm3d/hoho22k_2026_trainval split={args.split}...")
+    ds = load_dataset("usm3d/hoho22k_2026_trainval",
+                      streaming=True, trust_remote_code=True, split=args.split)
+
+    cfg = FuserConfig(depth_points_per_view=args.depth_per_view)
+    saved, skipped = 0, 0
+    t0 = time.perf_counter()
+    for i, sample in enumerate(ds):
+        if args.limit > 0 and i >= args.limit:
+            break
+        oid = sample["order_id"]
+        if oid in existing:
+            skipped += 1
+            continue
+        result = _process_one(sample, cfg)
+        if result is None:
+            skipped += 1
+            continue
+        order_id, data = result
+        torch.save(data, out_dir / f"{order_id}.pt")
+        saved += 1
+        if saved % 100 == 0:
+            rate = saved / (time.perf_counter() - t0)
+            print(f"  saved {saved} (skipped {skipped}) [{rate:.1f}/s]")
+
+    elapsed = time.perf_counter() - t0
+    print(f"Done. Saved {saved}, skipped {skipped} in {elapsed:.0f}s.")
+
+
+if __name__ == "__main__":
+    main()
+
+

s23dr_2026_example/color_mappings.py

@@ -0,0 +1,183 @@
+gestalt_color_mapping = {
+    "unclassified": (215, 62, 138),
+    "apex": (235, 88, 48),
+    "eave_end_point": (248, 130, 228),
+    "flashing_end_point": (71, 11, 161),
+    "ridge": (214, 251, 248),
+    "rake": (13, 94, 47),
+    "eave": (54, 243, 63),
+    "post": (187, 123, 236),
+    "ground_line": (136, 206, 14),
+    "flashing": (162, 162, 32),
+    "step_flashing": (169, 255, 219),
+    "hip": (8, 89, 52),
+    "valley": (85, 27, 65),
+    "roof": (215, 232, 179),
+    "door": (110, 52, 23),
+    "garage": (50, 233, 171),
+    "window": (230, 249, 40),
+    "shutter": (122, 4, 233),
+    "fascia": (95, 230, 240),
+    "soffit": (2, 102, 197),
+    "horizontal_siding": (131, 88, 59),
+    "vertical_siding": (110, 187, 198),
+    "brick": (171, 252, 7),
+    "concrete": (32, 47, 246),
+    "other_wall": (112, 61, 240),
+    "trim": (151, 206, 58),
+    "unknown": (127, 127, 127),
+    "transition_line": (0,0,0),
+}
+
+ade20k_color_mapping = {
+    'wall': (120, 120, 120),
+    'building;edifice': (180, 120, 120),
+    'sky': (6, 230, 230),
+    'floor;flooring': (80, 50, 50),
+    'tree': (4, 200, 3),
+    'ceiling': (120, 120, 80),
+    'road;route': (140, 140, 140),
+    'bed': (204, 5, 255),
+    'windowpane;window': (230, 230, 230),
+    'grass': (4, 250, 7),
+    'cabinet': (224, 5, 255),
+    'sidewalk;pavement': (235, 255, 7),
+    'person;individual;someone;somebody;mortal;soul': (150, 5, 61),
+    'earth;ground': (120, 120, 70),
+    'door;double;door': (8, 255, 51),
+    'table': (255, 6, 82),
+    'mountain;mount': (143, 255, 140),
+    'plant;flora;plant;life': (204, 255, 4),
+    'curtain;drape;drapery;mantle;pall': (255, 51, 7),
+    'chair': (204, 70, 3),
+    'car;auto;automobile;machine;motorcar': (0, 102, 200),
+    'water': (61, 230, 250),
+    'painting;picture': (255, 6, 51),
+    'sofa;couch;lounge': (11, 102, 255),
+    'shelf': (255, 7, 71),
+    'house': (255, 9, 224),
+    'sea': (9, 7, 230),
+    'mirror': (220, 220, 220),
+    'rug;carpet;carpeting': (255, 9, 92),
+    'field': (112, 9, 255),
+    'armchair': (8, 255, 214),
+    'seat': (7, 255, 224),
+    'fence;fencing': (255, 184, 6),
+    'desk': (10, 255, 71),
+    'rock;stone': (255, 41, 10),
+    'wardrobe;closet;press': (7, 255, 255),
+    'lamp': (224, 255, 8),
+    'bathtub;bathing;tub;bath;tub': (102, 8, 255),
+    'railing;rail': (255, 61, 6),
+    'cushion': (255, 194, 7),
+    'base;pedestal;stand': (255, 122, 8),
+    'box': (0, 255, 20),
+    'column;pillar': (255, 8, 41),
+    'signboard;sign': (255, 5, 153),
+    'chest;of;drawers;chest;bureau;dresser': (6, 51, 255),
+    'counter': (235, 12, 255),
+    'sand': (160, 150, 20),
+    'sink': (0, 163, 255),
+    'skyscraper': (140, 140, 140),
+    'fireplace;hearth;open;fireplace': (250, 10, 15),
+    'refrigerator;icebox': (20, 255, 0),
+    'grandstand;covered;stand': (31, 255, 0),
+    'path': (255, 31, 0),
+    'stairs;steps': (255, 224, 0),
+    'runway': (153, 255, 0),
+    'case;display;case;showcase;vitrine': (0, 0, 255),
+    'pool;table;billiard;table;snooker;table': (255, 71, 0),
+    'pillow': (0, 235, 255),
+    'screen;door;screen': (0, 173, 255),
+    'stairway;staircase': (31, 0, 255),
+    'river': (11, 200, 200),
+    'bridge;span': (255 ,82, 0),
+    'bookcase': (0, 255, 245),
+    'blind;screen': (0, 61, 255),
+    'coffee;table;cocktail;table': (0, 255, 112),
+    'toilet;can;commode;crapper;pot;potty;stool;throne': (0, 255, 133),
+    'flower': (255, 0, 0),
+    'book': (255, 163, 0),
+    'hill': (255, 102, 0),
+    'bench': (194, 255, 0),
+    'countertop': (0, 143, 255),
+    'stove;kitchen;stove;range;kitchen;range;cooking;stove': (51, 255, 0),
+    'palm;palm;tree': (0, 82, 255),
+    'kitchen;island': (0, 255, 41),
+    'computer;computing;machine;computing;device;data;processor;electronic;computer;information;processing;system': (0, 255, 173),
+    'swivel;chair': (10, 0, 255),
+    'boat': (173, 255, 0),
+    'bar': (0, 255, 153),
+    'arcade;machine': (255, 92, 0),
+    'hovel;hut;hutch;shack;shanty': (255, 0, 255),
+    'bus;autobus;coach;charabanc;double-decker;jitney;motorbus;motorcoach;omnibus;passenger;vehicle': (255, 0, 245),
+    'towel': (255, 0, 102),
+    'light;light;source': (255, 173, 0),
+    'truck;motortruck': (255, 0, 20),
+    'tower': (255, 184, 184),
+    'chandelier;pendant;pendent': (0, 31, 255),
+    'awning;sunshade;sunblind': (0, 255, 61),
+    'streetlight;street;lamp': (0, 71, 255),
+    'booth;cubicle;stall;kiosk': (255, 0, 204),
+    'television;television;receiver;television;set;tv;tv;set;idiot;box;boob;tube;telly;goggle;box': (0, 255, 194),
+    'airplane;aeroplane;plane': (0, 255, 82),
+    'dirt;track': (0, 10, 255),
+    'apparel;wearing;apparel;dress;clothes': (0, 112, 255),
+    'pole': (51, 0, 255),
+    'land;ground;soil': (0, 194, 255),
+    'bannister;banister;balustrade;balusters;handrail': (0, 122, 255),
+    'escalator;moving;staircase;moving;stairway': (0, 255, 163),
+    'ottoman;pouf;pouffe;puff;hassock': (255, 153, 0),
+    'bottle': (0, 255, 10),
+    'buffet;counter;sideboard': (255, 112, 0),
+    'poster;posting;placard;notice;bill;card': (143, 255, 0),
+    'stage': (82, 0, 255),
+    'van': (163, 255, 0),
+    'ship': (255, 235, 0),
+    'fountain': (8, 184, 170),
+    'conveyer;belt;conveyor;belt;conveyer;conveyor;transporter': (133, 0, 255),
+    'canopy': (0, 255, 92),
+    'washer;automatic;washer;washing;machine': (184, 0, 255),
+    'plaything;toy': (255, 0, 31),
+    'swimming;pool;swimming;bath;natatorium': (0, 184, 255),
+    'stool': (0, 214, 255),
+    'barrel;cask': (255, 0, 112),
+    'basket;handbasket': (92, 255, 0),
+    'waterfall;falls': (0, 224, 255),
+    'tent;collapsible;shelter': (112, 224, 255),
+    'bag': (70, 184, 160),
+    'minibike;motorbike': (163, 0, 255),
+    'cradle': (153, 0, 255),
+    'oven': (71, 255, 0),
+    'ball': (255, 0, 163),
+    'food;solid;food': (255, 204, 0),
+    'step;stair': (255, 0, 143),
+    'tank;storage;tank': (0, 255, 235),
+    'trade;name;brand;name;brand;marque': (133, 255, 0),
+    'microwave;microwave;oven': (255, 0, 235),
+    'pot;flowerpot': (245, 0, 255),
+    'animal;animate;being;beast;brute;creature;fauna': (255, 0, 122),
+    'bicycle;bike;wheel;cycle': (255, 245, 0),
+    'lake': (10, 190, 212),
+    'dishwasher;dish;washer;dishwashing;machine': (214, 255, 0),
+    'screen;silver;screen;projection;screen': (0, 204, 255),
+    'blanket;cover': (20, 0, 255),
+    'sculpture': (255, 255, 0),
+    'hood;exhaust;hood': (0, 153, 255),
+    'sconce': (0, 41, 255),
+    'vase': (0, 255, 204),
+    'traffic;light;traffic;signal;stoplight': (41, 0, 255),
+    'tray': (41, 255, 0),
+    'ashcan;trash;can;garbage;can;wastebin;ash;bin;ash-bin;ashbin;dustbin;trash;barrel;trash;bin': (173, 0, 255),
+    'fan': (0, 245, 255),
+    'pier;wharf;wharfage;dock': (71, 0, 255),
+    'crt;screen': (122, 0, 255),
+    'plate': (0, 255, 184),
+    'monitor;monitoring;device': (0, 92, 255),
+    'bulletin;board;notice;board': (184, 255, 0),
+    'shower': (0, 133, 255),
+    'radiator': (255, 214, 0),
+    'glass;drinking;glass': (25, 194, 194),
+    'clock': (102, 255, 0),
+    'flag': (92, 0, 255),
+}

s23dr_2026_example/data.py

@@ -0,0 +1,230 @@
+"""Data loading for pre-sampled HF datasets.
+
+Expects pre-sampled npz blobs with xyz_norm (not full PCD).
+Supports both 2048-point and 4096-point datasets.
+Use make_sampled_cache.py to produce these from full point clouds.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .tokenizer import EdgeDepthSequenceConfig
+
+# Default token budget (for 2048-point datasets; 4096 uses 3072/1024)
+SEQ_LEN = 2048
+COLMAP_POINTS = 1536
+DEPTH_POINTS = 512
+
+
+# ---------------------------------------------------------------------------
+# Datasets
+# ---------------------------------------------------------------------------
+
+def _load_bad_sample_ids():
+    """Load the set of known-bad sample IDs (misaligned GT, extreme scale)."""
+    bad_file = Path(__file__).parent / "bad_samples.txt"
+    if not bad_file.exists():
+        return set()
+    return set(line.strip() for line in bad_file.read_text().splitlines() if line.strip())
+
+
+class HFCachedDataset(torch.utils.data.Dataset):
+    """Load pre-sampled HuggingFace dataset into memory."""
+
+    def __init__(self, hf_dataset, aug_rotate=False, aug_jitter=0.0,
+                 aug_drop=0.0, aug_flip=False):
+        import io as _io
+        bad_ids = _load_bad_sample_ids()
+        print(f"Pre-decoding {len(hf_dataset)} samples into memory...")
+        self.samples = []
+        self.order_ids = []
+        n_skipped = 0
+        for i, sample in enumerate(hf_dataset):
+            if sample["order_id"] in bad_ids:
+                n_skipped += 1
+                continue
+            d = dict(np.load(_io.BytesIO(sample["data"])))
+            if "xyz_norm" not in d:
+                raise ValueError(
+                    f"Sample {sample['order_id']} missing 'xyz_norm' -- this looks like "
+                    f"a full PCD dataset, not pre-sampled. Use make_sampled_cache.py first.")
+            self.samples.append(d)
+            self.order_ids.append(sample["order_id"])
+            if (i + 1) % 2000 == 0:
+                print(f"  {i+1}/{len(hf_dataset)}...")
+        print(f"  Done. {len(self.samples)} samples in memory"
+              f" ({n_skipped} bad samples filtered).")
+        self.aug_rotate = aug_rotate
+        self.aug_jitter = aug_jitter
+        self.aug_drop = aug_drop
+        self.aug_flip = aug_flip
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        out = _process_sample(self.samples[idx], self.aug_rotate,
+                              self.aug_jitter, self.aug_drop, self.aug_flip)
+        out["sample_id"] = self.order_ids[idx]
+        return out
+
+
+def _process_sample(d, aug_rotate, aug_jitter=0.0, aug_drop=0.0, aug_flip=False):
+    """Process a pre-sampled npz dict into training tensors.
+
+    Args:
+        aug_rotate: random yaw rotation
+        aug_jitter: std of Gaussian noise added to point positions (0=disabled)
+        aug_drop: fraction of points to randomly drop (0=disabled)
+        aug_flip: random mirror along X axis (50% chance)
+    """
+    xyz_norm = d["xyz_norm"].copy()
+    gt_seg = d["gt_segments"].copy()
+    mask = d["mask"].copy()
+
+    if aug_rotate:
+        theta = np.random.rand() * 2 * np.pi
+        cos_t, sin_t = np.cos(theta), np.sin(theta)
+        x, z = xyz_norm[:, 0].copy(), xyz_norm[:, 2].copy()
+        xyz_norm[:, 0] = x * cos_t - z * sin_t
+        xyz_norm[:, 2] = x * sin_t + z * cos_t
+        for ep in range(2):
+            sx, sz = gt_seg[:, ep, 0].copy(), gt_seg[:, ep, 2].copy()
+            gt_seg[:, ep, 0] = sx * cos_t - sz * sin_t
+            gt_seg[:, ep, 2] = sx * sin_t + sz * cos_t
+
+    if aug_flip and np.random.rand() < 0.5:
+        xyz_norm[:, 0] = -xyz_norm[:, 0]
+        gt_seg[:, :, 0] = -gt_seg[:, :, 0]
+
+    if aug_jitter > 0:
+        valid = mask.astype(bool)
+        xyz_norm[valid] += np.random.randn(valid.sum(), 3).astype(np.float32) * aug_jitter
+
+    if aug_drop > 0:
+        valid_idx = np.where(mask)[0]
+        n_drop = int(len(valid_idx) * aug_drop)
+        if n_drop > 0:
+            drop_idx = np.random.choice(valid_idx, n_drop, replace=False)
+            mask[drop_idx] = False
+
+    result = {
+        "xyz_norm": torch.as_tensor(xyz_norm, dtype=torch.float32),
+        "class_id": torch.as_tensor(d["class_id"], dtype=torch.long),
+        "source": torch.as_tensor(d["source"], dtype=torch.long),
+        "mask": torch.as_tensor(mask),
+        "gt_segments": torch.as_tensor(gt_seg, dtype=torch.float32),
+        "scale": torch.tensor(float(d["scale"]), dtype=torch.float32),
+        "center": torch.as_tensor(d["center"], dtype=torch.float32),
+        "gt_vertices": d["gt_vertices"],
+        "gt_edges": d["gt_edges"],
+        "visible_src": torch.as_tensor(d["visible_src"], dtype=torch.long),
+        "visible_id": torch.as_tensor(d["visible_id"], dtype=torch.long),
+    }
+    if "behind" in d:
+        result["behind"] = torch.as_tensor(
+            np.clip(np.asarray(d["behind"], dtype=np.int16), 0, None), dtype=torch.long)
+    if "n_views_voted" in d:
+        result["n_views_voted"] = torch.as_tensor(d["n_views_voted"], dtype=torch.float32)
+    if "vote_frac" in d:
+        result["vote_frac"] = torch.as_tensor(d["vote_frac"], dtype=torch.float32)
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Collation + DataLoader
+# ---------------------------------------------------------------------------
+
+def collate(batch):
+    """Stack samples into batched tensors."""
+    out = {
+        "xyz_norm": torch.stack([d["xyz_norm"] for d in batch]),
+        "class_id": torch.stack([d["class_id"] for d in batch]),
+        "source": torch.stack([d["source"] for d in batch]),
+        "mask": torch.stack([d["mask"] for d in batch]),
+        "gt_segments": [d["gt_segments"] for d in batch],
+        "scales": torch.stack([d["scale"] for d in batch]),
+        "meta": batch,
+    }
+    # Optional fields: check ALL samples, not just batch[0].
+    # If any sample has it, all must have it (no mixed data versions).
+    for field in ("behind", "n_views_voted", "vote_frac"):
+        if any(field in d for d in batch):
+            missing = [i for i, d in enumerate(batch) if field not in d]
+            if missing:
+                raise KeyError(
+                    f"Field '{field}' present in some batch samples but missing in "
+                    f"{len(missing)}/{len(batch)}. Mixed data versions in cache?")
+            out[field] = torch.stack([d[field] for d in batch])
+    return out
+
+
+def build_loader(cache_dir, batch_size, aug_rotate=False, aug_jitter=0.0,
+                 aug_drop=0.0, aug_flip=False):
+    """Create a DataLoader from HF dataset.
+
+    cache_dir should be 'hf://repo/name:split' format.
+    """
+    if cache_dir.startswith("local://"):
+        from datasets import load_from_disk
+        hf_ds = load_from_disk(cache_dir[len("local://"):])
+    elif cache_dir.startswith("hf://"):
+        parts = cache_dir[5:].split(":")
+        repo = parts[0]
+        split = parts[1] if len(parts) > 1 else "train"
+        from datasets import load_dataset
+        hf_ds = load_dataset(repo, split=split)
+    else:
+        raise ValueError(
+            f"cache_dir must be 'hf://repo:split' or 'local:///path', got: {cache_dir}")
+    ds = HFCachedDataset(hf_ds, aug_rotate=aug_rotate, aug_jitter=aug_jitter,
+                          aug_drop=aug_drop, aug_flip=aug_flip)
+    loader = torch.utils.data.DataLoader(
+        ds, batch_size=batch_size, shuffle=True,
+        num_workers=0, collate_fn=collate,
+    )
+    print(f"Dataset: {len(ds)} scenes, batch_size={batch_size}")
+    return loader
+
+
+# ---------------------------------------------------------------------------
+# Token building (GPU)
+# ---------------------------------------------------------------------------
+
+def build_tokens(batch, model, device):
+    """Apply Fourier features + learned embeddings on GPU."""
+    xyz = batch["xyz_norm"].to(device)
+    cid = batch["class_id"].to(device)
+    src = batch["source"].to(device)
+    masks = batch["mask"].to(device)
+    gt = [g.to(device) for g in batch["gt_segments"]]
+    scales = batch["scales"]
+
+    B, T, _ = xyz.shape
+    tok = model.tokenizer
+    fourier = tok.pos_enc(xyz.reshape(-1, 3)).reshape(B, T, -1) \
+        if tok.pos_enc is not None else xyz.new_zeros(B, T, 0)
+    parts = [xyz, fourier, tok.label_emb(cid), tok.src_emb(src.clamp(0, 1))]
+    if tok.behind_emb_dim > 0:
+        if "behind" in batch:
+            beh = batch["behind"].to(device)
+        else:
+            # Data doesn't have behind -- use zeros (embed index 0).
+            # This is intentional for eval on old data; for training,
+            # fail fast by requiring the field (checked in _process_sample).
+            beh = xyz.new_zeros(B, T, dtype=torch.long)
+        parts.append(tok.behind_emb(beh))
+    if tok.use_vote_features:
+        if "n_views_voted" not in batch or "vote_frac" not in batch:
+            raise KeyError(
+                "Model expects vote features (--vote-features) but data is missing "
+                "'n_views_voted'/'vote_frac'. Use v2 dataset or regenerate cache.")
+        # Normalize to ~zero mean, unit variance (dataset stats: nv~2.7+/-1.0, vf~0.5+/-0.25)
+        nv = ((batch["n_views_voted"].to(device).float() - 2.7) / 1.0).unsqueeze(-1)
+        vf = ((batch["vote_frac"].to(device).float() - 0.5) / 0.25).unsqueeze(-1)
+        parts.extend([nv, vf])
+    tokens = torch.cat(parts, dim=-1)
+    return tokens, masks, gt, scales, batch["meta"]

s23dr_2026_example/losses.py

@@ -0,0 +1,215 @@
+"""Loss computation for wireframe prediction."""
+from __future__ import annotations
+
+import torch
+
+from .varifold import varifold_loss_batch
+from .sinkhorn import batched_sinkhorn_loss
+
+# Varifold config
+VARIANT = "simpson3"
+SIGMAS = [0.5, 1.0, 2.0]  # meters (divided by per-scene scale at runtime)
+ALPHAS = [0.2, 0.6, 0.2]
+LEN_POW = 1.0
+VARIFOLD_CROSS_ONLY = False  # Set to True to drop self-energy (avoids O(S^2) blowup)
+
+# Sinkhorn config (note: near-zero gradients at eps=0.05, effectively disabled)
+SINKHORN_EPS = 0.05
+SINKHORN_ITERS = 10
+
+# Sinkhorn dustbin cost: controls the OT "not matching" penalty.
+# Like tau, this is an OT behavior parameter, NOT a physical distance.
+# Must be comparable to typical matching costs in normalized space (~0.1).
+# Do NOT divide by scale.
+SINKHORN_DUSTBIN = 0.1
+
+MAX_GT = 64  # fixed pad size for compile-friendly shapes
+
+# Precomputed constants (created once on first call)
+_loss_constants = {}
+
+
+def _get_loss_constants(device, dtype):
+    key = (device, dtype)
+    if key not in _loss_constants:
+        _loss_constants[key] = {
+            "sigmas": torch.tensor(SIGMAS, device=device, dtype=dtype),
+            "alphas": torch.tensor(ALPHAS, device=device, dtype=dtype),
+        }
+    return _loss_constants[key]
+
+
+def pad_gt_fixed(gt_list, device, dtype):
+    """Pad GT segments to fixed MAX_GT for compile-friendly shapes."""
+    B = len(gt_list)
+    gt_pad = torch.zeros((B, MAX_GT, 2, 3), device=device, dtype=dtype)
+    gt_mask = torch.zeros((B, MAX_GT), device=device, dtype=torch.bool)
+    gt_lengths = torch.zeros(B, device=device, dtype=dtype)
+    for i, g in enumerate(gt_list):
+        n = g.shape[0]
+        if n > 0:
+            gt_pad[i, :n] = g
+            gt_mask[i, :n] = True
+            gt_lengths[i] = torch.linalg.norm(g[:, 1] - g[:, 0], dim=-1).sum()
+    return gt_pad, gt_mask, gt_lengths
+
+
+def _loss_inner(pred_segments, gt_pad, gt_mask, gt_lengths, scales,
+                sigmas, alphas, varifold_w):
+    """Pure tensor loss -- no Python control flow, no boolean indexing."""
+    has_gt = (gt_lengths > 0).float()
+
+    sigmas_eff = sigmas / scales[:, None]
+    loss_batch = varifold_loss_batch(
+        pred_segments, gt_pad, gt_mask=gt_mask,
+        variant=VARIANT, sigmas=sigmas_eff, alpha=alphas, len_pow=LEN_POW,
+        cross_only=VARIFOLD_CROSS_ONLY,
+    )
+    v = loss_batch / gt_lengths.clamp(min=1.0)
+    v = (v * has_gt).sum() / has_gt.sum().clamp(min=1.0)
+
+    total = varifold_w * v
+    return total, v
+
+
+# Will be replaced with compiled version on CUDA
+_loss_fn = _loss_inner
+
+
+def compute_loss(pred_segments, gt_list, scales, device,
+                 varifold_w, sinkhorn_w,
+                 endpoint_w=0.0,
+                 conf_logits=None, conf_weight=0.0, conf_mode="sinkhorn",
+                 sinkhorn_eps=None, sinkhorn_iters=None,
+                 sinkhorn_dustbin=None, conf_clamp_min=None):
+    """Combined loss with fixed-size GT padding.
+
+    conf_mode: "sinkhorn" = conf-weighted sinkhorn, "sinkhorn_detach" = detached conf.
+    """
+    if conf_logits is not None and conf_clamp_min is not None:
+        conf_logits = conf_logits.clamp(min=conf_clamp_min)
+    gt_pad, gt_mask, gt_lengths = pad_gt_fixed(gt_list, device, pred_segments.dtype)
+    c = _get_loss_constants(device, pred_segments.dtype)
+
+    total, v = _loss_fn(
+        pred_segments, gt_pad, gt_mask, gt_lengths, scales,
+        c["sigmas"], c["alphas"], varifold_w)
+
+    terms = {}
+    if varifold_w > 0:
+        terms["varifold"] = v.detach()
+
+    if sinkhorn_w > 0:
+        has_gt = (gt_lengths > 0).float()
+        if conf_logits is not None and conf_mode == "sinkhorn":
+            pred_mass = torch.sigmoid(conf_logits)
+        elif conf_logits is not None and conf_mode == "sinkhorn_detach":
+            pred_mass = torch.sigmoid(conf_logits.detach())
+        else:
+            pred_mass = None
+        eps = sinkhorn_eps if sinkhorn_eps is not None else SINKHORN_EPS
+        iters = sinkhorn_iters if sinkhorn_iters is not None else SINKHORN_ITERS
+        dustbin = sinkhorn_dustbin if sinkhorn_dustbin is not None else SINKHORN_DUSTBIN
+        S = pred_segments.shape[1]
+        sink_per = batched_sinkhorn_loss(
+            pred_segments, gt_pad, gt_mask,
+            eps, iters, dustbin,
+            pred_mass=pred_mass,
+        ) / (gt_lengths.clamp(min=1.0) * S)
+        s = (sink_per * has_gt).sum() / has_gt.sum().clamp(min=1.0)
+        total = total + sinkhorn_w * s
+        terms["sinkhorn"] = s.detach()
+
+    if conf_logits is not None and conf_weight > 0:
+        if conf_mode in ("sinkhorn", "sinkhorn_detach"):
+            conf_w = torch.sigmoid(conf_logits)
+            S = conf_logits.shape[1]
+            gt_counts = gt_mask.sum(dim=1).float()
+            conf_sum = conf_w.sum(dim=1)
+            reg = (((conf_sum - gt_counts) / S) ** 2).mean()
+            total = total + conf_weight * reg
+            terms["conf_reg"] = reg.detach()
+        else:
+            raise ValueError(f"Unknown conf_mode: {conf_mode}")
+
+    if endpoint_w > 0:
+        has_gt = (gt_lengths > 0).float()
+        eps_ep = sinkhorn_eps if sinkhorn_eps is not None else SINKHORN_EPS
+        iters_ep = sinkhorn_iters if sinkhorn_iters is not None else SINKHORN_ITERS
+        dustbin_ep = sinkhorn_dustbin if sinkhorn_dustbin is not None else SINKHORN_DUSTBIN
+        B, S = pred_segments.shape[:2]
+        M = gt_pad.shape[1]
+
+        # Compute hard assignment via sinkhorn (detached -- matching is not trained)
+        with torch.no_grad():
+            pred_mass_ep = torch.sigmoid(conf_logits) if conf_logits is not None else None
+            sink_loss_for_assign = batched_sinkhorn_loss(
+                pred_segments, gt_pad, gt_mask, eps_ep, iters_ep, dustbin_ep,
+                pred_mass=pred_mass_ep)
+            p0, p1 = pred_segments[:, :, 0], pred_segments[:, :, 1]
+            g0, g1 = gt_pad[:, :, 0], gt_pad[:, :, 1]
+            mid_p, half_p = 0.5 * (p0 + p1), 0.5 * (p1 - p0)
+            mid_g, half_g = 0.5 * (g0 + g1), 0.5 * (g1 - g0)
+            d_mid = torch.linalg.norm(mid_p.unsqueeze(2) - mid_g.unsqueeze(1), dim=-1)
+            len_p = torch.linalg.norm(half_p, dim=-1, keepdim=True).clamp(min=1e-6)
+            len_g = torch.linalg.norm(half_g, dim=-1, keepdim=True).clamp(min=1e-6)
+            dir_p, dir_g = half_p / len_p, half_g / len_g
+            cos_a = (dir_p.unsqueeze(2) * dir_g.unsqueeze(1)).sum(dim=-1)
+            d_dir = 1.0 - cos_a.abs()
+            d_len = (len_p.unsqueeze(2) - len_g.unsqueeze(1)).squeeze(-1).abs()
+            cost = d_mid + d_dir + d_len
+            dc = torch.as_tensor(dustbin_ep, device=cost.device, dtype=cost.dtype)
+            cost = torch.where(gt_mask.unsqueeze(1), cost, dc * 10.0)
+            cost_pad = dc.expand(B, S + 1, M + 1).clone()
+            cost_pad[:, :S, :M] = cost
+            cost_pad[:, -1, -1] = 0.0
+            gt_counts = gt_mask.sum(dim=1).float()
+            if pred_mass_ep is not None:
+                pm = pred_mass_ep.clamp(min=0.0)
+                a = torch.cat([pm, (gt_counts - pm.sum(1)).clamp(min=0).unsqueeze(1)], dim=1)
+                b_val = torch.zeros(B, M + 1, device=cost.device, dtype=cost.dtype)
+                b_val[:, :M] = gt_mask.float()
+                b_val[:, -1] = (pm.sum(1) - gt_counts).clamp(min=0)
+            else:
+                n = float(S)
+                denom = n + gt_counts
+                a = (1.0 / denom).unsqueeze(1).expand(B, S + 1).clone()
+                a[:, -1] = gt_counts / denom
+                b_val = (1.0 / denom).unsqueeze(1).expand(B, M + 1).clone()
+                b_val[:, -1] = n / denom
+                b_val[:, :M] = b_val[:, :M] * gt_mask.float()
+            log_a = torch.log(a + 1e-9)
+            log_b = torch.log(b_val + 1e-9)
+            log_k = -cost_pad / eps_ep
+            log_u = torch.zeros_like(a)
+            log_v = torch.zeros_like(b_val)
+            for _ in range(iters_ep):
+                log_u = log_a - torch.logsumexp(log_k + log_v.unsqueeze(1), dim=2)
+                log_v = log_b - torch.logsumexp(log_k + log_u.unsqueeze(2), dim=1)
+            transport = torch.exp(log_u.unsqueeze(2) + log_v.unsqueeze(1) + log_k)
+            assignment = transport[:, :S, :M+1].argmax(dim=2)
+            assignment[assignment >= M] = -1
+
+        # Everything below is WITH gradients (assignment is detached but pred_segments is live)
+        matched = (assignment >= 0)  # [B, S]
+        n_matched = matched.float().sum().clamp(min=1.0)
+        assign_safe = assignment.clamp(min=0)
+        gt_matched = gt_pad[
+            torch.arange(B, device=device)[:, None].expand(B, S),
+            assign_safe]  # [B, S, 2, 3]
+
+        # Symmetric endpoint distance
+        ref_ep1 = pred_segments[:, :, 0]
+        ref_ep2 = pred_segments[:, :, 1]
+        gt_ep1 = gt_matched[:, :, 0]
+        gt_ep2 = gt_matched[:, :, 1]
+        dist_fwd = (ref_ep1 - gt_ep1).norm(dim=-1) + (ref_ep2 - gt_ep2).norm(dim=-1)
+        dist_rev = (ref_ep1 - gt_ep2).norm(dim=-1) + (ref_ep2 - gt_ep1).norm(dim=-1)
+        ep_dist = torch.min(dist_fwd, dist_rev)
+
+        # Normalize by GT total length * S (same scale as sinkhorn)
+        ep_loss = (ep_dist * matched.float()).sum() / n_matched
+        total = total + endpoint_w * ep_loss
+        terms["endpoint"] = ep_loss.detach()
+
+    return total, terms

s23dr_2026_example/make_sampled_cache.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python3
+"""Stage 2: priority-sample cached .pt scenes into fixed-size .npz files.
+
+Reads the per-scene .pt files produced by cache_scenes.py, priority-samples
+a fixed number of points (2048 or 4096), normalizes, and writes one .npz per
+scene (~50KB at 2048, ~100KB at 4096).
+
+A fixed seed is used so every scene gets one deterministic sample -- no
+per-epoch sampling augmentation, every epoch sees the same points.
+
+Usage:
+  python -m s23dr_2026_example.make_sampled_cache \\
+      --in-dir cache/full --out-dir cache/sampled_2048 --seq-len 2048
+  python -m s23dr_2026_example.make_sampled_cache \\
+      --in-dir cache/full --out-dir cache/sampled_4096 --seq-len 4096
+
+The 3:1 colmap:depth quota ratio is fixed: at seq_len=2048 that's
+colmap=1536/depth=512; at seq_len=4096 that's colmap=3072/depth=1024.
+"""
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+# Priority sampling (same logic as train.py)
+def _priority_sample(source, group_id, seq_len, colmap_quota, depth_quota):
+    def pick(src_id, quota):
+        base = source == src_id
+        picked, remaining = [], quota
+        for tier in range(5):
+            if remaining <= 0:
+                break
+            pool = np.where(base & (group_id == tier))[0]
+            if len(pool) == 0:
+                continue
+            np.random.shuffle(pool)
+            take = min(remaining, len(pool))
+            picked.append(pool[:take])
+            remaining -= take
+        if remaining > 0:
+            pool = np.where(base & (group_id >= 0))[0]
+            if len(pool) > 0:
+                np.random.shuffle(pool)
+                picked.append(pool[:min(remaining, len(pool))])
+                remaining -= min(remaining, len(pool))
+        return np.concatenate(picked) if picked else np.array([], dtype=np.int64), remaining
+
+    idx_c, rem_c = pick(0, colmap_quota)
+    idx_d, rem_d = pick(1, depth_quota)
+
+    if rem_c > 0:
+        extra = np.setdiff1d(np.where((source == 1) & (group_id >= 0))[0], idx_d)
+        np.random.shuffle(extra)
+        idx_d = np.concatenate([idx_d, extra[:rem_c]])
+    if rem_d > 0:
+        extra = np.setdiff1d(np.where((source == 0) & (group_id >= 0))[0], idx_c)
+        np.random.shuffle(extra)
+        idx_c = np.concatenate([idx_c, extra[:rem_d]])
+
+    indices = np.concatenate([idx_c, idx_d])
+    num_valid = len(indices)
+    if num_valid < seq_len:
+        if num_valid == 0:
+            return np.zeros(seq_len, dtype=np.int64), np.zeros(seq_len, dtype=bool)
+        indices = np.concatenate([indices, np.full(seq_len - num_valid, indices[-1])])
+    mask = np.zeros(seq_len, dtype=bool)
+    mask[:num_valid] = True
+    return indices[:seq_len], mask
+
+
+def _process_sample(d, seq_len, colmap_q, depth_q):
+    """Sample and normalize one cached scene dict into a small npz-ready dict."""
+    xyz = np.asarray(d["xyz"], np.float32)
+    source = np.asarray(d["source"], np.uint8)
+    group_id = np.asarray(d["group_id"], np.int8)
+    class_id = np.asarray(d["class_id"], np.uint8)
+    vis_src = np.asarray(d["visible_src"], np.uint8)
+    vis_id = np.asarray(d["visible_id"], np.int16)
+    center = np.asarray(d["center"], np.float32)
+    scale = float(d["scale"])
+    gt_v = np.asarray(d["gt_vertices"], np.float32)
+    gt_e = np.asarray(d["gt_edges"], np.int32)
+
+    indices, mask = _priority_sample(source, group_id, seq_len, colmap_q, depth_q)
+    xyz_norm = ((xyz[indices] - center) / scale).astype(np.float32)
+    gt_seg = np.stack([gt_v[gt_e[:, 0]], gt_v[gt_e[:, 1]]], axis=1)
+    gt_seg_norm = ((gt_seg - center) / scale).astype(np.float32)
+
+    result = {
+        "xyz_norm": xyz_norm,
+        "class_id": class_id[indices].astype(np.uint8),
+        "source": source[indices].astype(np.uint8),
+        "mask": mask,
+        "gt_segments": gt_seg_norm,
+        "scale": np.float32(scale),
+        "center": center,
+        "gt_vertices": gt_v,
+        "gt_edges": gt_e,
+        "visible_src": vis_src[indices].astype(np.uint8),
+        "visible_id": vis_id[indices].astype(np.int16),
+    }
+    if "behind_gest_id" in d:
+        result["behind"] = np.asarray(d["behind_gest_id"], np.int16)[indices]
+    if "n_views_voted" in d:
+        result["n_views_voted"] = np.asarray(d["n_views_voted"], np.uint8)[indices]
+    if "vote_frac" in d:
+        result["vote_frac"] = np.asarray(d["vote_frac"], np.float32)[indices]
+    if "gt_edge_classes" in d:
+        result["gt_edge_classes"] = np.asarray(d["gt_edge_classes"], np.int64)
+    return result
+
+
+def main():
+    p = argparse.ArgumentParser(description="Stage 2: cached .pt -> sampled .npz")
+    p.add_argument("--in-dir", required=True, help="Directory of .pt files from cache_scenes.py")
+    p.add_argument("--out-dir", required=True, help="Output directory for .npz files")
+    p.add_argument("--seq-len", type=int, default=2048, help="Points per sample (2048 or 4096)")
+    p.add_argument("--seed", type=int, default=7)
+    args = p.parse_args()
+
+    colmap_q = args.seq_len * 3 // 4
+    depth_q = args.seq_len - colmap_q
+    print(f"seq_len={args.seq_len}  colmap={colmap_q}  depth={depth_q}")
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    np.random.seed(args.seed)
+
+    files = sorted(Path(args.in_dir).glob("*.pt"))
+    print(f"Found {len(files)} .pt files in {args.in_dir}")
+
+    done = 0
+    t0 = time.perf_counter()
+    for f in files:
+        out_f = out_dir / (f.stem + ".npz")
+        if out_f.exists():
+            done += 1
+            continue
+        d = torch.load(f, weights_only=False)
+        result = _process_sample(d, args.seq_len, colmap_q, depth_q)
+        np.savez(out_f, **result)
+        done += 1
+        if done % 2000 == 0:
+            rate = done / (time.perf_counter() - t0)
+            print(f"  {done}/{len(files)} [{rate:.0f}/s]")
+
+    elapsed = time.perf_counter() - t0
+    print(f"Done. {done} files in {elapsed:.0f}s -> {out_dir}")
+
+
+if __name__ == "__main__":
+    main()
+
+

s23dr_2026_example/model.py

@@ -0,0 +1,519 @@
+"""
+Perceiver-based transformer for 3D roof wireframe prediction.
+
+Architecture overview:
+
+    Input tokens [B, T, D]
+        |
+        v
+    input_proj: Linear -> GELU -> Linear -> LayerNorm   =>  [B, T, hidden]
+        |
+        v
+    Perceiver latent bottleneck (N PerceiverLatentLayers):
+        Learnable latent embeddings [L, hidden] are broadcast to batch.
+        Each layer: cross-attn(latents <- tokens) -> self-attn(latents) -> FFN
+        Output: latents [B, L, hidden]
+        |
+        v
+    Segment decoder (M SegmentDecoderLayers):
+        Learnable query embeddings [S, hidden] are broadcast to batch.
+        Each layer: cross-attn(queries <- latents) -> self-attn(queries) -> FFN
+        Output: queries [B, S, hidden]
+        |
+        v
+    segment_head: Linear -> 6D -> (midpoint, half_vector)
+        + query_offsets (learnable per-query bias)
+        endpoints = midpoint +/- half_vector  ->  [B, S, 2, 3]
+"""
+
+import torch
+import torch.nn as nn
+
+from .attention import MultiHeadSDPA, FeedForward
+
+
+# ---------------------------------------------------------------------------
+# Building blocks
+# ---------------------------------------------------------------------------
+
+class AttnResidual(nn.Module):
+    """Pre-norm attention + residual + dropout."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        kv_heads: int | None = None,
+        norm_class=None,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        norm_class = norm_class or nn.LayerNorm
+        self.norm = norm_class(d_model)
+        self.attn = MultiHeadSDPA(d_model, num_heads, kv_heads=kv_heads, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        memory: torch.Tensor,
+        memory_key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        res = x
+        x = self.norm(x)
+        x = self.attn(x, memory, key_padding_mask=memory_key_padding_mask)
+        return res + self.drop(x)
+
+
+class FFNResidual(nn.Module):
+    """Pre-norm feed-forward + residual + dropout."""
+
+    def __init__(
+        self,
+        d_model: int,
+        dim_ff: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        norm_class=None,
+    ):
+        super().__init__()
+        norm_class = norm_class or nn.LayerNorm
+        self.norm = norm_class(d_model)
+        self.ffn = FeedForward(d_model, dim_ff, activation=activation)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        res = x
+        x = self.norm(x)
+        x = self.ffn(x)
+        return res + self.drop(x)
+
+
+# ---------------------------------------------------------------------------
+# Perceiver encoder layer
+# ---------------------------------------------------------------------------
+
+class PerceiverLatentLayer(nn.Module):
+    """Single Perceiver latent layer.
+
+    If use_cross=True:  cross-attn(latents <- points) -> self-attn -> FFN
+    If use_cross=False: self-attn -> FFN  (saves compute in deep stacks)
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_ff: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        kv_heads_cross: int | None = None,
+        kv_heads_self: int | None = None,
+        use_cross: bool = True,
+        norm_class=None,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        self.use_cross = use_cross
+        if use_cross:
+            self.cross = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_self, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)
+
+    def forward(
+        self,
+        latents: torch.Tensor,
+        points: torch.Tensor,
+        points_key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.use_cross:
+            latents = self.cross(latents, points, memory_key_padding_mask=points_key_padding_mask)
+        latents = self.self_attn(latents, latents)
+        latents = self.ffn(latents)
+        return latents
+
+
+# ---------------------------------------------------------------------------
+# Segment decoder layer
+# ---------------------------------------------------------------------------
+
+class SegmentDecoderLayer(nn.Module):
+    """Single segment decoder layer.
+
+    cross-attn(queries <- latents) -> [cross-attn(queries <- inputs)] -> self-attn(queries) -> FFN
+
+    If input_xattn=True, adds a second cross-attention that attends directly
+    to the projected input tokens (bypassing the latent bottleneck). This gives
+    queries access to fine-grained point-level detail for vertex precision.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_ff: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        kv_heads_cross: int | None = None,
+        kv_heads_self: int | None = None,
+        norm_class=None,
+        input_xattn: bool = False,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        self.cross = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.input_xattn = input_xattn
+        if input_xattn:
+            self.cross_input = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_cross, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads_self, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)
+
+    def forward(
+        self,
+        queries: torch.Tensor,
+        latents: torch.Tensor,
+        src: torch.Tensor | None = None,
+        src_key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        queries = self.cross(queries, latents)
+        if self.input_xattn and src is not None:
+            queries = self.cross_input(queries, src, memory_key_padding_mask=src_key_padding_mask)
+        queries = self.self_attn(queries, queries)
+        queries = self.ffn(queries)
+        return queries
+
+
+# ---------------------------------------------------------------------------
+# Full model
+# ---------------------------------------------------------------------------
+
+class TokenTransformerSegments(nn.Module):
+    """Perceiver transformer that predicts 3D roof wireframe segments.
+
+    Takes point-cloud tokens and outputs segment endpoints as [B, S, 2, 3]
+    where S is the number of segments and each segment has two 3D endpoints.
+
+    Args:
+        segments:       Number of predicted segments (S).
+        in_dim:         Dimensionality of input tokens.
+        hidden:         Internal hidden dimension throughout the model.
+        num_heads:      Number of attention heads.
+        kv_heads_cross: Grouped-query heads for cross-attention (None = standard MHA).
+        kv_heads_self:  Grouped-query heads for self-attention (None = standard MHA).
+        dim_feedforward: FFN intermediate dimension.
+        dropout:        Dropout rate applied after attention and FFN.
+        latent_tokens:  Number of learnable latent embeddings (L) in the bottleneck.
+        latent_layers:  Number of PerceiverLatentLayers (N).
+        decoder_layers: Number of SegmentDecoderLayers (M).
+    """
+
+    def __init__(
+        self,
+        segments: int = 32,
+        in_dim: int = 128,
+        hidden: int = 128,
+        num_heads: int = 4,
+        kv_heads_cross: int | None = 2,
+        kv_heads_self: int | None = 0,
+        dim_feedforward: int = 256,
+        dropout: float = 0.01,
+        latent_tokens: int = 64,
+        latent_layers: int = 2,
+        decoder_layers: int = 2,
+        cross_attn_interval: int = 1,
+        norm_class=None,
+        activation: str = "gelu",
+        segment_conf: bool = False,
+        pre_encoder_layers: int = 0,
+        segment_param: str = "midpoint_halfvec",
+        length_floor: float = 0.0,
+        decoder_input_xattn: bool = False,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        self.segments = segments
+        self.out_vertices = segments * 2
+        self.segment_param = segment_param
+        self.decoder_input_xattn = decoder_input_xattn
+        norm_class = norm_class or nn.LayerNorm
+
+        # Treat 0 as "use standard MHA"
+        if kv_heads_cross is not None and kv_heads_cross <= 0:
+            kv_heads_cross = None
+        if kv_heads_self is not None and kv_heads_self <= 0:
+            kv_heads_self = None
+
+        # -- Input projection --
+        self.input_proj = nn.Sequential(
+            nn.Linear(in_dim, dim_feedforward),
+            nn.GELU(),
+            nn.Linear(dim_feedforward, hidden),
+            norm_class(hidden),
+        )
+
+        # -- Optional pre-encoder: self-attention on full token sequence --
+        if pre_encoder_layers > 0:
+            self.pre_encoder = nn.ModuleList([
+                SelfAttentionEncoderLayer(
+                    d_model=hidden,
+                    num_heads=num_heads,
+                    dim_ff=dim_feedforward,
+                    dropout=dropout,
+                    activation=activation,
+                    kv_heads=kv_heads_self,
+                    norm_class=norm_class,
+                    qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+                )
+                for _ in range(pre_encoder_layers)
+            ])
+        else:
+            self.pre_encoder = None
+
+        # -- Perceiver latent bottleneck --
+        self.latent_embed = nn.Embedding(latent_tokens, hidden)
+        N = latent_layers
+        self.latent_layers = nn.ModuleList([
+            PerceiverLatentLayer(
+                d_model=hidden,
+                num_heads=num_heads,
+                dim_ff=dim_feedforward,
+                dropout=dropout,
+                activation=activation,
+                kv_heads_cross=kv_heads_cross,
+                kv_heads_self=kv_heads_self,
+                use_cross=(i == 0) or (i == N - 1) or (i % cross_attn_interval == 0),
+                norm_class=norm_class,
+                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+            )
+            for i in range(N)
+        ])
+
+        # -- Segment decoder --
+        self.query_embed = nn.Embedding(segments, hidden)
+        self.decoder_layers = nn.ModuleList([
+            SegmentDecoderLayer(
+                d_model=hidden,
+                num_heads=num_heads,
+                dim_ff=dim_feedforward,
+                dropout=dropout,
+                activation=activation,
+                kv_heads_cross=kv_heads_cross,
+                kv_heads_self=kv_heads_self,
+                norm_class=norm_class,
+                input_xattn=decoder_input_xattn,
+                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+            )
+            for _ in range(decoder_layers)
+        ])
+
+        # -- Output head --
+        if segment_param == "midpoint_dir_len":
+            self.segment_head = nn.Linear(hidden, 7)  # mid(3) + dir(3) + len(1)
+        else:
+            self.segment_head = nn.Linear(hidden, 6)  # mid(3) + half(3)
+        self.query_offsets = nn.Parameter(torch.zeros(segments, 2, 3))
+
+        nn.init.trunc_normal_(self.segment_head.weight, mean=0.0, std=1e-3)
+        if self.segment_head.bias is not None:
+            nn.init.zeros_(self.segment_head.bias)
+            if segment_param == "midpoint_dir_len":
+                # softplus(0.5) * 0.1 ~= 0.097 default length in normalized space
+                self.segment_head.bias.data[6] = 0.5
+        nn.init.normal_(self.query_offsets, mean=0.0, std=0.05)
+
+        # -- Optional confidence head --
+        self.segment_conf = segment_conf
+        if segment_conf:
+            self.conf_head = nn.Linear(hidden, 1)
+            nn.init.zeros_(self.conf_head.bias)
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        mask: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor | list]:
+        """
+        Args:
+            tokens: Input point-cloud tokens [B, T, in_dim].
+            mask:   Boolean validity mask [B, T]. True = valid token.
+
+        Returns:
+            Dict with keys:
+                "vertices": [B, S*2, 3] flattened endpoints.
+                "segments": [B, S, 2, 3] segment endpoints.
+                "edges":    Per-batch list of (start, end) index pairs into vertices.
+                "conf":     [B, S] logits (only if segment_conf=True).
+        """
+        B = tokens.shape[0]
+
+        # Project input tokens
+        src = self.input_proj(tokens)  # [B, T, hidden]
+
+        # Padding mask (True where padded) for cross-attention
+        pad_mask = ~mask.bool() if mask is not None else None
+
+        # Optional pre-encoder: self-attention on full token sequence
+        if self.pre_encoder is not None:
+            for layer in self.pre_encoder:
+                src = layer(src, key_padding_mask=pad_mask)
+
+        # Perceiver latent bottleneck
+        latents = self.latent_embed.weight.unsqueeze(0).expand(B, -1, -1)
+        for layer in self.latent_layers:
+            latents = layer(latents, src, points_key_padding_mask=pad_mask)
+
+        # Segment decoder
+        queries = self.query_embed.weight.unsqueeze(0).expand(B, -1, -1)
+        for layer in self.decoder_layers:
+            queries = layer(queries, latents,
+                            src=src if self.decoder_input_xattn else None,
+                            src_key_padding_mask=pad_mask if self.decoder_input_xattn else None)
+
+        # Predict segments -> endpoints
+        if self.segment_param == "midpoint_dir_len":
+            raw = self.segment_head(queries)  # [B, S, 7]
+            mid = raw[:, :, :3] + self.query_offsets[:, 0, :].unsqueeze(0)
+            direction = torch.nn.functional.normalize(raw[:, :, 3:6], dim=-1)
+            length = torch.nn.functional.softplus(raw[:, :, 6:7]) * 0.1
+            half = direction * length * 0.5
+        else:
+            raw = self.segment_head(queries).view(B, self.segments, 2, 3)
+            raw = raw + self.query_offsets.unsqueeze(0)
+            mid, half = raw[:, :, 0], raw[:, :, 1]
+        seg_params = torch.stack([mid - half, mid + half], dim=2)
+
+        vertices = seg_params.reshape(B, self.out_vertices, 3)
+        edges = [[(2 * i, 2 * i + 1) for i in range(self.segments)] for _ in range(B)]
+
+        out = {"vertices": vertices, "segments": seg_params, "edges": edges,
+               "src": src, "pad_mask": pad_mask, "queries": queries}
+        if self.segment_conf:
+            out["conf"] = self.conf_head(queries).squeeze(-1)  # [B, S]
+        return out
+
+
+# ---------------------------------------------------------------------------
+# Encoder-only layer (self-attention on full token sequence)
+# ---------------------------------------------------------------------------
+
+class SelfAttentionEncoderLayer(nn.Module):
+    """Single self-attention layer: self-attn(tokens) -> FFN."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_ff: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        kv_heads: int | None = None,
+        norm_class=None,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+    ):
+        super().__init__()
+        self.self_attn = AttnResidual(d_model, num_heads, dropout, kv_heads=kv_heads, norm_class=norm_class, qk_norm=qk_norm, qk_norm_type=qk_norm_type)
+        self.ffn = FFNResidual(d_model, dim_ff, dropout, activation=activation, norm_class=norm_class)
+
+    def forward(self, x: torch.Tensor, key_padding_mask: torch.Tensor | None = None) -> torch.Tensor:
+        x = self.self_attn(x, x, memory_key_padding_mask=key_padding_mask)
+        x = self.ffn(x)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# End-to-end model: tokenizer embeddings + perceiver
+# ---------------------------------------------------------------------------
+
+class EdgeDepthSegmentsModel(nn.Module):
+    """Tokenizer embeddings + transformer for 3D roof wireframes.
+
+    Supports two architectures via the `arch` parameter:
+    - "perceiver": Perceiver latent bottleneck (default, O(L*T) attention)
+    - "transformer": Standard self-attention encoder (O(T^2) attention)
+
+    Both share the same decoder, output head, and tokenizer.
+    """
+
+    def __init__(
+        self,
+        seq_cfg,
+        segments: int = 32,
+        hidden: int = 128,
+        num_heads: int = 4,
+        kv_heads_cross: int | None = 2,
+        kv_heads_self: int | None = 0,
+        dim_feedforward: int = 256,
+        dropout: float = 0.1,
+        latent_tokens: int = 64,
+        latent_layers: int = 1,
+        decoder_layers: int = 2,
+        label_emb_dim: int = 16,
+        src_emb_dim: int = 2,
+        behind_emb_dim: int = 8,
+        fourier_seed: int = 0,
+        cross_attn_interval: int = 1,
+        norm_class=None,
+        activation: str = "gelu",
+        segment_conf: bool = False,
+        use_vote_features: bool = False,
+        arch: str = "perceiver",
+        encoder_layers: int = 4,
+        pre_encoder_layers: int = 0,
+        segment_param: str = "midpoint_halfvec",
+        length_floor: float = 0.0,
+        decoder_input_xattn: bool = False,
+        qk_norm: bool = False,
+        qk_norm_type: str = "l2",
+        learnable_fourier: bool = False,
+    ):
+        super().__init__()
+        self.seq_cfg = seq_cfg
+
+        from .tokenizer import EdgeDepthSequenceBuilder
+        self.tokenizer = EdgeDepthSequenceBuilder(
+            seq_cfg,
+            label_emb_dim=label_emb_dim,
+            src_emb_dim=src_emb_dim,
+            behind_emb_dim=behind_emb_dim,
+            fourier_seed=fourier_seed,
+            use_vote_features=use_vote_features,
+            learnable_fourier=learnable_fourier,
+        )
+
+        if arch == "transformer":
+            raise ValueError(
+                "arch='transformer' is no longer supported. "
+                "TransformerSegments has been removed; use arch='perceiver'.")
+        else:
+            self.segmenter = TokenTransformerSegments(
+                segments=segments,
+                in_dim=self.tokenizer.out_dim,
+                hidden=hidden,
+                num_heads=num_heads,
+                kv_heads_cross=kv_heads_cross,
+                kv_heads_self=kv_heads_self,
+                dim_feedforward=dim_feedforward,
+                dropout=dropout,
+                latent_tokens=latent_tokens,
+                latent_layers=latent_layers,
+                decoder_layers=decoder_layers,
+                cross_attn_interval=cross_attn_interval,
+                norm_class=norm_class,
+                activation=activation,
+                segment_conf=segment_conf,
+                pre_encoder_layers=pre_encoder_layers,
+                segment_param=segment_param,
+                length_floor=length_floor,
+                decoder_input_xattn=decoder_input_xattn,
+                qk_norm=qk_norm, qk_norm_type=qk_norm_type,
+            )
+
+    def forward_tokens(self, tokens: torch.Tensor, mask: torch.Tensor):
+        """Run the segmenter on pre-built token tensors."""
+        return self.segmenter(tokens, mask)

s23dr_2026_example/point_fusion.py

@@ -0,0 +1,554 @@
+"""
+point_fusion.py
+
+Simplified semantic point fusion for the 2026 dataset format.
+
+Takes per-view (ADE segmap, Gestalt segmap, depth) + sparse COLMAP point cloud
+from the usm3d/hoho22k_2026_trainval dataset and builds a compact, house-centric
+semantic point representation suitable for downstream wireframe prediction.
+
+Key differences from the 2025 pipeline:
+  - COLMAP is a ZIP of text files (cameras.txt, images.txt, points3D.txt)
+  - Depth is millimeter I;16 PNG (depth_scale=0.001 converts to meters)
+  - Views flagged with pose_only_in_colmap=True have zeroed K/R/t and must be
+    skipped for depth unprojection and projection
+  - Images arrive as PIL Images, not byte arrays
+"""
+
+from __future__ import annotations
+
+import zipfile
+from dataclasses import dataclass
+from io import BytesIO
+from typing import Dict, List, Optional, Tuple
+
+import cv2
+import numpy as np
+from scipy.stats import mode as scipy_mode
+
+from .color_mappings import ade20k_color_mapping, gestalt_color_mapping
+
+# ---------------------------------------------------------------------------
+# Color packing helpers
+# ---------------------------------------------------------------------------
+
+def _pack_rgb_u32(rgb: np.ndarray) -> np.ndarray:
+    """Pack uint8 RGB (..., 3) into uint32 codes."""
+    rgb = rgb.astype(np.uint32, copy=False)
+    return (rgb[..., 0] << 16) | (rgb[..., 1] << 8) | rgb[..., 2]
+
+
+def _build_rgbcode_maps(color_mapping):
+    """Return (rgbcode_to_id, id_to_name) for a color mapping dict."""
+    names = list(color_mapping.keys())
+    rgbs = np.array([color_mapping[n] for n in names], dtype=np.uint8)
+    codes = _pack_rgb_u32(rgbs.reshape(-1, 1, 3)).reshape(-1)
+    rgbcode_to_id = {int(c): i for i, c in enumerate(codes)}
+    return rgbcode_to_id, names
+
+
+def _name_to_packed_rgb(name, mapping):
+    """Case-insensitive lookup returning a packed RGB code, or None."""
+    for key in mapping:
+        if key.lower() == name.lower():
+            rgb = np.array(mapping[key], np.uint8).reshape(1, 1, 3)
+            return int(_pack_rgb_u32(rgb).reshape(()))
+    return None
+
+# ---------------------------------------------------------------------------
+# Label mapping constants
+# ---------------------------------------------------------------------------
+
+ADE_RGBCODE_TO_ID, ADE_ID_TO_NAME = _build_rgbcode_maps(ade20k_color_mapping)
+GEST_RGBCODE_TO_ID, GEST_ID_TO_NAME = _build_rgbcode_maps(gestalt_color_mapping)
+NUM_ADE = len(ADE_ID_TO_NAME)
+NUM_GEST = len(GEST_ID_TO_NAME)
+
+GEST_INVALID_NAMES = ("unclassified", "unknown", "transition_line")
+GEST_INVALID_CODES = set(
+    int(_pack_rgb_u32(np.array(gestalt_color_mapping[n], np.uint8).reshape(1, 1, 3)).reshape(()))
+    for n in GEST_INVALID_NAMES if n in gestalt_color_mapping
+)
+
+# ADE classes whose surfaces are "see-through" for label fusion: when a point
+# projects onto one of these, we use the Gestalt label behind it instead.
+ADE_TRANSPARENT_NAMES = (
+    "wall", "building;edifice", "floor;flooring", "ceiling",
+    "windowpane;window", "door;double;door", "house", "skyscraper",
+    "screen;door;screen", "blind;screen", "hovel;hut;hutch;shack;shanty",
+    "tower", "booth;cubicle;stall;kiosk",
+)
+
+# ADE classes kept as "occluders/add-ons" when overlapping the house silhouette.
+ADE_OCCLUDER_ALLOWLIST_NAMES = (
+    "tree", "person;individual;someone;somebody;mortal;soul",
+    "car;auto;automobile;machine;motorcar", "truck;motortruck", "van",
+    "fence;fencing", "railing;rail",
+    "bannister;banister;balustrade;balusters;handrail",
+    "stairs;steps", "stairway;staircase", "step;stair", "pole",
+    "streetlight;street;lamp", "signboard;sign", "awning;sunshade;sunblind",
+    "plant;flora;plant;life", "pot;flowerpot",
+)
+
+# Precomputed arrays for the default name lists (avoids re-lookup every call).
+_DEFAULT_ADE_TRANSPARENT_CODES = np.array(
+    [c for n in ADE_TRANSPARENT_NAMES
+     if (c := _name_to_packed_rgb(n, ade20k_color_mapping)) is not None],
+    dtype=np.uint32,
+)
+_DEFAULT_ADE_OCCLUDER_IDS = np.array(
+    sorted({ADE_RGBCODE_TO_ID[c]
+            for n in ADE_OCCLUDER_ALLOWLIST_NAMES
+            if (c := _name_to_packed_rgb(n, ade20k_color_mapping)) is not None
+            and c in ADE_RGBCODE_TO_ID}),
+    dtype=np.int32,
+)
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+@dataclass(frozen=True)
+class FuserConfig:
+    """Simplified fusion configuration (no depth calibration fields)."""
+    depth_points_per_view: int = 20_000          # depth samples per view
+    depth_scale: float = 0.001                   # mm -> meters
+    depth_clip_percentile: float = 99.5          # drop extreme outliers
+    house_mask_dilate_px: int = 5                # dilate gestalt mask
+    min_support_views: int = 1                   # min views for a kept point
+    ade_transparent_classes: Tuple[str, ...] = ADE_TRANSPARENT_NAMES
+    ade_occluder_allowlist: Tuple[str, ...] = ADE_OCCLUDER_ALLOWLIST_NAMES
+
+# ---------------------------------------------------------------------------
+# Geometry: projection + depth unprojection
+# ---------------------------------------------------------------------------
+
+def project_world_points(points_world, K, R, t):
+    """Project (N,3) world points to pixel (u,v) with validity mask."""
+    pts = points_world.astype(np.float32, copy=False)
+    cam = (R @ pts.T + t).T  # (N, 3)
+    z = cam[:, 2]
+    valid = z > 1e-6
+    inv_z = np.zeros_like(z)
+    inv_z[valid] = 1.0 / z[valid]
+    x = cam[:, 0] * inv_z
+    y = cam[:, 1] * inv_z
+    u = K[0, 0] * x + K[0, 2]
+    v = K[1, 1] * y + K[1, 2]
+    return u, v, valid
+
+
+def unproject_depth_to_world(depth, K, R, t, num_points, sample_mask=None, rng=None):
+    """Convert a depth map + camera params to (M, 3) world points, M <= num_points."""
+    if rng is None:
+        rng = np.random.default_rng()
+    d = np.asarray(depth, dtype=np.float32)
+    if d.ndim != 2:
+        return np.zeros((0, 3), dtype=np.float32)
+
+    valid = np.isfinite(d) & (d > 1e-6)
+    if sample_mask is not None:
+        mask = np.asarray(sample_mask, dtype=bool)
+        if mask.shape != d.shape:
+            return np.zeros((0, 3), dtype=np.float32)
+        valid &= mask
+
+    ys, xs = np.where(valid)
+    if ys.size == 0:
+        return np.zeros((0, 3), dtype=np.float32)
+
+    idx = rng.choice(ys.size, size=min(num_points, ys.size), replace=False)
+    y = ys[idx].astype(np.float32)
+    x = xs[idx].astype(np.float32)
+    z = d[ys[idx], xs[idx]].astype(np.float32)
+
+    fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
+    cam_pts = np.stack([(x - cx) * z / fx, (y - cy) * z / fy, z], axis=0)
+    # cam = R * world + t  =>  world = R^T * (cam - t)
+    world = (R.T @ (cam_pts - t)).T
+    return world.astype(np.float32, copy=False)
+
+
+def clean_depth(depth, clip_percentile):
+    """Clip extreme depth values."""
+    d = np.asarray(depth, dtype=np.float32)
+    d = np.where(np.isfinite(d), d, 0.0)
+    d[d <= 0] = 0.0
+    if clip_percentile is not None and clip_percentile > 0 and np.any(d > 0):
+        hi = float(np.percentile(d[d > 0], clip_percentile))
+        d = np.clip(d, 0.0, hi)
+    return d
+
+
+def dilate_mask(mask, radius_px):
+    """Binary dilation via cv2.  mask: (H, W) bool."""
+    if radius_px <= 0:
+        return mask
+    k = 2 * radius_px + 1
+    kernel = np.ones((k, k), np.uint8)
+    return cv2.dilate(mask.astype(np.uint8), kernel) > 0
+
+# ---------------------------------------------------------------------------
+# COLMAP extraction (2026 format)
+# ---------------------------------------------------------------------------
+
+def extract_colmap_points_2026(sample):
+    """Extract (N, 3) float32 COLMAP world points from a 2026-format sample.
+
+    sample['colmap'] must be a ZIP archive containing points3D.txt.
+    Fails fast if that file is missing (it is always present in the 2026 format).
+    """
+    colmap_blob = sample.get("colmap")
+    if colmap_blob is None:
+        return np.zeros((0, 3), dtype=np.float32)
+    if not isinstance(colmap_blob, (bytes, bytearray, memoryview)):
+        return np.zeros((0, 3), dtype=np.float32)
+
+    try:
+        with zipfile.ZipFile(BytesIO(colmap_blob)) as zf:
+            if "points3D.txt" not in set(zf.namelist()):
+                raise FileNotFoundError(
+                    "COLMAP ZIP is missing points3D.txt -- "
+                    "this is required in the 2026 dataset format")
+            with zf.open("points3D.txt") as f:
+                text = f.read().decode("utf-8", errors="ignore")
+            # Format: POINT3D_ID X Y Z R G B ERROR TRACK[]
+            # Filter comment/blank lines, parse columns 1-3 (X,Y,Z)
+            from io import StringIO
+            clean = "\n".join(l for l in text.split("\n") if l and not l.startswith("#"))
+            if not clean:
+                return np.zeros((0, 3), dtype=np.float32)
+            return np.loadtxt(StringIO(clean), dtype=np.float32, usecols=(1, 2, 3))
+    except zipfile.BadZipFile:
+        pass
+    return np.zeros((0, 3), dtype=np.float32)
+
+# ---------------------------------------------------------------------------
+# Label helpers
+# ---------------------------------------------------------------------------
+
+def _codes_from_image(img):
+    """Convert a PIL Image or numpy array to a (H, W) uint32 packed-RGB map."""
+    arr = np.asarray(img)
+    if arr.ndim == 2:
+        arr = np.stack([arr, arr, arr], axis=-1)
+    arr = arr[..., :3]
+    if arr.dtype != np.uint8:
+        arr = np.clip(arr, 0, 255).astype(np.uint8)
+    return _pack_rgb_u32(arr)
+
+
+def _row_majority(values):
+    """Row-wise majority vote on (P, V) int array; -1 means "no vote".
+    Returns (P,) with the most frequent non-negative value per row, or -1.
+
+    Masks -1 entries before voting so that abstentions don't outvote
+    actual labels (which happens when a point is visible in only 1-2 views).
+    """
+    P, V = values.shape
+    result = np.full(P, -1, dtype=values.dtype)
+
+    # For each row, find the most frequent non-negative value.
+    # Vectorized approach: flatten valid entries per row using argmax on counts.
+    # Since values are typically small non-negative ints (0-200), we can use
+    # a simple max-of-first-valid approach for speed when V is small.
+    for vi in range(V):
+        # For rows still unset, take the first valid vote
+        col = values[:, vi]
+        unset = result == -1
+        has_val = col >= 0
+        update = unset & has_val
+        result[update] = col[update]
+
+    # Now refine: if a row has multiple different valid votes, pick the mode.
+    # Check if any row has conflicting votes across views.
+    has_any = np.any(values >= 0, axis=1)
+    n_valid = np.sum(values >= 0, axis=1)
+    needs_vote = has_any & (n_valid > 1)
+
+    if np.any(needs_vote):
+        for i in np.where(needs_vote)[0]:
+            valid = values[i][values[i] >= 0]
+            # Use numpy bincount for speed (values are small non-neg ints)
+            counts = np.bincount(valid.astype(np.intp))
+            result[i] = counts.argmax()
+
+    return result
+
+# ---------------------------------------------------------------------------
+# Semantic fusion: house-centric, occluder-aware
+# ---------------------------------------------------------------------------
+
+def _fuse_labels_for_points(
+    points_world, Ks, Rs, ts, ade_images, gestalt_images,
+    ade_transparent_codes, ade_occluder_allowed_ids,
+    min_support_views, valid_view_mask=None,
+):
+    """Multi-view semantic label fusion with majority voting.
+
+    For each 3D point, project into every valid view:
+      - ADE "envelope" class -> use the Gestalt label behind it.
+      - ADE non-envelope -> keep if on the occluder allowlist.
+    Then majority-vote across views.
+
+    Returns dict: keep, visible_src, visible_id, behind_gest_id, support
+    """
+    P = points_world.shape[0]
+    V = min(len(Ks), len(Rs), len(ts), len(ade_images), len(gestalt_images))
+    empty = {
+        "keep": np.zeros(P, dtype=bool),
+        "visible_src": np.zeros(P, np.uint8),
+        "visible_id": np.full(P, -1, np.int16),
+        "behind_gest_id": np.full(P, -1, np.int16),
+        "support": np.zeros(P, np.uint8),
+    }
+    if P == 0 or V == 0:
+        return empty
+
+    # Per-view labels. src: 1=gestalt, 2=ade; -1 = no contribution.
+    visible_src_pv = np.full((P, V), -1, dtype=np.int8)
+    visible_id_pv = np.full((P, V), -1, dtype=np.int32)
+    behind_id_pv = np.full((P, V), -1, dtype=np.int32)
+    support = np.zeros(P, dtype=np.int32)
+
+    ade_allowed_set = set(ade_occluder_allowed_ids.tolist())
+    ade_transparent_u32 = ade_transparent_codes.astype(np.uint32, copy=False)
+    gest_invalid_arr = np.array(list(GEST_INVALID_CODES), dtype=np.uint32)
+
+    for vi in range(V):
+        if valid_view_mask is not None and not valid_view_mask[vi]:
+            continue
+
+        K = np.asarray(Ks[vi], np.float32)
+        R = np.asarray(Rs[vi], np.float32)
+        t = np.asarray(ts[vi], np.float32).reshape(3, 1)
+
+        ade_codes_img = _codes_from_image(ade_images[vi])
+        gest_codes_img = _codes_from_image(gestalt_images[vi])
+        H, W = ade_codes_img.shape
+
+        u, v, valid = project_world_points(points_world, K, R, t)
+        in_img = valid & (u >= 0) & (u < W) & (v >= 0) & (v < H)
+        if not np.any(in_img):
+            continue
+
+        ui = np.clip(np.round(u[in_img]).astype(np.int32), 0, W - 1)
+        vi_pix = np.clip(np.round(v[in_img]).astype(np.int32), 0, H - 1)
+        ade_codes = ade_codes_img[vi_pix, ui]
+        gest_codes = gest_codes_img[vi_pix, ui]
+
+        in_house = ~np.isin(gest_codes, gest_invalid_arr)
+        if not np.any(in_house):
+            continue
+
+        idx = np.where(in_img)[0][in_house]
+        ade_codes_h = ade_codes[in_house]
+        gest_codes_h = gest_codes[in_house]
+
+        behind_local = np.array(
+            [GEST_RGBCODE_TO_ID.get(int(c), -1) for c in gest_codes_h],
+            dtype=np.int32)
+        behind_id_pv[idx, vi] = behind_local
+
+        ade_is_transparent = np.isin(ade_codes_h, ade_transparent_u32)
+
+        # Case A: ADE is envelope -- use Gestalt label.
+        mask_a = ade_is_transparent & (behind_local >= 0)
+        if np.any(mask_a):
+            visible_src_pv[idx[mask_a], vi] = 1
+            visible_id_pv[idx[mask_a], vi] = behind_local[mask_a]
+
+        # Case B: ADE is non-envelope -- use ADE label (allowlist-filtered).
+        mask_b = ~ade_is_transparent
+        if np.any(mask_b):
+            ade_local = np.array(
+                [ADE_RGBCODE_TO_ID.get(int(c), -1) for c in ade_codes_h[mask_b]],
+                dtype=np.int32)
+            on_allowlist = np.array(
+                [int(a) in ade_allowed_set for a in ade_local], dtype=bool
+            ) & (ade_local >= 0)
+            if np.any(on_allowlist):
+                visible_src_pv[idx[mask_b][on_allowlist], vi] = 2
+                visible_id_pv[idx[mask_b][on_allowlist], vi] = ade_local[on_allowlist]
+
+        support[idx] += 1
+
+    # ---- Aggregate across views via majority vote ----
+    keep = (support >= min_support_views) & np.any(visible_src_pv >= 0, axis=1)
+
+    # Combine (src, id) into a single key for voting, then split back.
+    # src in {1,2} and id in [0, ~150], so stride=100k avoids collisions.
+    VIS_STRIDE = 100_000
+    vis_key = np.where(
+        visible_src_pv >= 0,
+        visible_src_pv.astype(np.int64) * VIS_STRIDE + visible_id_pv.astype(np.int64),
+        -1)
+    voted_key = _row_majority(vis_key)
+    voted_behind = _row_majority(behind_id_pv)
+
+    final_src = np.zeros(P, dtype=np.uint8)
+    final_id = np.full(P, -1, dtype=np.int16)
+    ok = voted_key >= 0
+    if np.any(ok):
+        final_src[ok] = (voted_key[ok] // VIS_STRIDE).astype(np.uint8)
+        final_id[ok] = (voted_key[ok] % VIS_STRIDE).astype(np.int16)
+
+    # ---- Vote confidence metadata ----
+    n_views_voted = np.sum(visible_src_pv >= 0, axis=1).astype(np.uint8)
+
+    # Fraction of voting views that agreed with the majority label
+    vote_frac = np.zeros(P, dtype=np.float32)
+    if np.any(ok):
+        for i in np.where(ok)[0]:
+            votes = vis_key[i][vis_key[i] >= 0]
+            if len(votes) > 0:
+                vote_frac[i] = (votes == voted_key[i]).sum() / len(votes)
+
+    return {
+        "keep": keep,
+        "visible_src": final_src,
+        "visible_id": final_id,
+        "behind_gest_id": voted_behind.astype(np.int16),
+        "support": support.astype(np.uint8),
+        "n_views_voted": n_views_voted,
+        "vote_frac": vote_frac,
+    }
+
+# ---------------------------------------------------------------------------
+# Compact scene builder (2026 dataset format)
+# ---------------------------------------------------------------------------
+
+def _resolve_ade_codes(cfg):
+    """Return (transparent_codes, occluder_ids) for the given config.
+    Uses precomputed module-level arrays when the config has default names.
+    """
+    if cfg.ade_transparent_classes == ADE_TRANSPARENT_NAMES:
+        transparent = _DEFAULT_ADE_TRANSPARENT_CODES
+    else:
+        transparent = np.array(
+            [c for n in cfg.ade_transparent_classes
+             if (c := _name_to_packed_rgb(n, ade20k_color_mapping)) is not None],
+            dtype=np.uint32)
+
+    if cfg.ade_occluder_allowlist == ADE_OCCLUDER_ALLOWLIST_NAMES:
+        occluder_ids = _DEFAULT_ADE_OCCLUDER_IDS
+    else:
+        occluder_ids = np.array(
+            sorted({ADE_RGBCODE_TO_ID[c]
+                    for n in cfg.ade_occluder_allowlist
+                    if (c := _name_to_packed_rgb(n, ade20k_color_mapping)) is not None
+                    and c in ADE_RGBCODE_TO_ID}),
+            dtype=np.int32)
+    return transparent, occluder_ids
+
+
+def _parse_gt_array(sample, key, dtype, expected_cols):
+    """Parse an optional ground-truth array from the sample dict."""
+    raw = sample.get(key)
+    if raw is None:
+        return None
+    arr = np.asarray(raw, dtype=dtype)
+    if arr.ndim == 2 and arr.shape[1] == expected_cols:
+        return arr
+    return None
+
+
+def build_compact_scene(sample, cfg, rng):
+    """Build a compact semantic point representation from a HuggingFace sample.
+
+    Expected sample keys: K, R, t, ade, gestalt, depth, colmap,
+    pose_only_in_colmap, wf_vertices (opt), wf_edges (opt), __key__ (opt).
+
+    Returns dict (xyz, source, visible_src, visible_id, behind_gest_id,
+    gt_vertices, gt_edges, sample_id) or None if no points survive fusion.
+    """
+    Ks = sample.get("K") or []
+    Rs = sample.get("R") or []
+    ts = sample.get("t") or []
+    ade_imgs = sample.get("ade") or []
+    gest_imgs = sample.get("gestalt") or []
+    depths = sample.get("depth") or []
+    pose_flags = sample.get("pose_only_in_colmap") or []
+
+    V = min(len(Ks), len(Rs), len(ts), len(ade_imgs), len(gest_imgs))
+    if V == 0:
+        return None
+
+    valid_view = [not (vi < len(pose_flags) and pose_flags[vi]) for vi in range(V)]
+    if not any(valid_view):
+        return None
+
+    # ---- COLMAP points ----
+    colmap_pts = extract_colmap_points_2026(sample)
+
+    # ---- Precompute house masks (from Gestalt), optionally dilated ----
+    gest_invalid_arr = np.array(list(GEST_INVALID_CODES), dtype=np.uint32)
+    house_masks = []
+    for vi in range(V):
+        if not valid_view[vi]:
+            house_masks.append(None)
+            continue
+        mask = ~np.isin(_codes_from_image(gest_imgs[vi]), gest_invalid_arr)
+        if cfg.house_mask_dilate_px > 0:
+            mask = dilate_mask(mask, cfg.house_mask_dilate_px)
+        house_masks.append(mask)
+
+    # ---- Sample depth points per view ----
+    depth_points_all = []
+    for vi in range(min(V, len(depths))):
+        if not valid_view[vi] or depths[vi] is None:
+            continue
+        d = clean_depth(
+            np.asarray(depths[vi], dtype=np.float32) * cfg.depth_scale,
+            cfg.depth_clip_percentile)
+        pts = unproject_depth_to_world(
+            depth=d,
+            K=np.asarray(Ks[vi], np.float32),
+            R=np.asarray(Rs[vi], np.float32),
+            t=np.asarray(ts[vi], np.float32).reshape(3, 1),
+            num_points=cfg.depth_points_per_view,
+            sample_mask=house_masks[vi], rng=rng)
+        if pts.shape[0]:
+            depth_points_all.append(pts)
+
+    # ---- Combine COLMAP + depth points ----
+    pts_list, src_list = [], []
+    if colmap_pts.shape[0]:
+        pts_list.append(colmap_pts)
+        src_list.append(np.zeros(colmap_pts.shape[0], dtype=np.uint8))   # 0=colmap
+    if depth_points_all:
+        all_depth = np.concatenate(depth_points_all, axis=0)
+        pts_list.append(all_depth)
+        src_list.append(np.ones(all_depth.shape[0], dtype=np.uint8))     # 1=depth
+    if not pts_list:
+        return None
+
+    points_world = np.concatenate(pts_list, axis=0).astype(np.float32, copy=False)
+    point_source = np.concatenate(src_list, axis=0).astype(np.uint8, copy=False)
+
+    # ---- Fuse semantic labels ----
+    ade_transparent_arr, ade_allow_ids = _resolve_ade_codes(cfg)
+    fused = _fuse_labels_for_points(
+        points_world=points_world, Ks=Ks, Rs=Rs, ts=ts,
+        ade_images=ade_imgs, gestalt_images=gest_imgs,
+        ade_transparent_codes=ade_transparent_arr,
+        ade_occluder_allowed_ids=ade_allow_ids,
+        min_support_views=cfg.min_support_views,
+        valid_view_mask=valid_view)
+
+    keep = fused["keep"]
+    if not np.any(keep):
+        return None
+
+    return {
+        "xyz": points_world[keep],
+        "source": point_source[keep],               # 0=colmap, 1=monodepth
+        "visible_src": fused["visible_src"][keep],   # 1=gestalt, 2=ade
+        "visible_id": fused["visible_id"][keep],
+        "behind_gest_id": fused["behind_gest_id"][keep],
+        "n_views_voted": fused["n_views_voted"][keep],
+        "vote_frac": fused["vote_frac"][keep],
+        "gt_vertices": _parse_gt_array(sample, "wf_vertices", np.float32, 3),
+        "gt_edges": _parse_gt_array(sample, "wf_edges", np.int64, 2),
+        "sample_id": sample.get("__key__", None),
+    }

s23dr_2026_example/postprocess_v2.py

@@ -0,0 +1,39 @@
+"""Post-processing functions for segment predictions."""
+import numpy as np
+
+
+def snap_to_point_cloud(vertices, xyz, class_id, snap_radius=0.5,
+                         target_classes=None):
+    """Snap vertices to nearby point cloud clusters of specific semantic classes."""
+    if target_classes is None:
+        target_classes = [1, 2]  # apex, eave_end_point
+
+    snapped = vertices.copy()
+    mask = np.isin(class_id, target_classes)
+
+    if mask.sum() < 2:
+        return snapped
+
+    target_pts = xyz[mask]
+
+    for i, v in enumerate(vertices):
+        dists = np.linalg.norm(target_pts - v, axis=-1)
+        close = dists < snap_radius
+        if close.sum() >= 2:
+            snapped[i] = target_pts[close].mean(axis=0)
+
+    return snapped
+
+
+def snap_horizontal(vertices, edges, max_slope=0.05):
+    """Snap near-horizontal edges to be exactly horizontal."""
+    verts = vertices.copy()
+    for a, b in edges:
+        a, b = int(a), int(b)
+        dy = abs(verts[a, 1] - verts[b, 1])
+        dxz = np.sqrt((verts[a, 0] - verts[b, 0])**2 + (verts[a, 2] - verts[b, 2])**2)
+        if dxz > 0.1 and dy / dxz < max_slope:
+            avg_y = 0.5 * (verts[a, 1] + verts[b, 1])
+            verts[a, 1] = avg_y
+            verts[b, 1] = avg_y
+    return verts

s23dr_2026_example/segment_postprocess.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+def merge_vertices_iterative(vertices: np.ndarray, edges: np.ndarray,
+                              start: float = 0.15, end: float = 0.6,
+                              n_iters: int = 5):
+    """Iterative merge: start with tight threshold, gradually widen.
+
+    Avoids the worst transitive chaining effects of a single wide threshold.
+    Each pass merges only the closest pairs first, establishing stable cluster
+    centers before wider merges pull in more distant endpoints.
+
+    +0.004 HSS / +0.007 F1 over single-pass merge(0.4) on 1024 val samples.
+    """
+    pv, pe = vertices, edges
+    for t in np.linspace(start, end, n_iters):
+        pv, pe = merge_vertices(pv, pe, t)
+    return pv, pe
+
+
+def merge_vertices(vertices: np.ndarray, edges: np.ndarray, thresh: float):
+    verts = np.asarray(vertices, dtype=np.float32)
+    edges = np.asarray(edges, dtype=np.int64)
+    if verts.size == 0 or edges.size == 0:
+        return verts, edges
+
+    n = verts.shape[0]
+    parent = np.arange(n, dtype=np.int64)
+
+    def find(i):
+        while parent[i] != i:
+            parent[i] = parent[parent[i]]
+            i = parent[i]
+        return i
+
+    def union(i, j):
+        ri = find(i)
+        rj = find(j)
+        if ri != rj:
+            parent[rj] = ri
+
+    for i in range(n):
+        vi = verts[i]
+        for j in range(i + 1, n):
+            if np.linalg.norm(vi - verts[j]) <= thresh:
+                union(i, j)
+
+    clusters = {}
+    for i in range(n):
+        root = find(i)
+        clusters.setdefault(root, []).append(i)
+
+    new_vertices = []
+    mapping = {}
+    for new_idx, idxs in enumerate(clusters.values()):
+        pts = verts[idxs]
+        center = pts.mean(axis=0)
+        new_vertices.append(center)
+        for i in idxs:
+            mapping[i] = new_idx
+
+    new_edges = []
+    seen = set()
+    for a, b in edges:
+        na = mapping.get(int(a), int(a))
+        nb = mapping.get(int(b), int(b))
+        if na == nb:
+            continue
+        key = (na, nb) if na <= nb else (nb, na)
+        if key in seen:
+            continue
+        seen.add(key)
+        new_edges.append([na, nb])
+
+    return np.asarray(new_vertices, dtype=np.float32), np.asarray(new_edges, dtype=np.int64)

s23dr_2026_example/sinkhorn.py

@@ -0,0 +1,126 @@
+"""Sinkhorn optimal transport loss for segment matching.
+
+Note: at eps=0.05, sinkhorn gradients are near-zero (~1e-7 norm) for
+typical matrix sizes. The loss value is tracked but does not meaningfully
+train the model. Default sinkhorn_weight=0.0. See worklog.md for details.
+
+Future: schedule eps from large (1.0) to small (0.05) during training
+to get useful gradients early and precise matching late.
+"""
+import torch
+
+
+def batched_sinkhorn_loss(
+    pred_segments: torch.Tensor,
+    gt_pad: torch.Tensor,
+    gt_mask: torch.Tensor,
+    eps: float,
+    iters: int,
+    dustbin_cost: float | torch.Tensor,
+    pred_mass: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Batched sinkhorn segment matching loss.
+
+    Args:
+        pred_segments: [B, S, 2, 3] predicted segments
+        gt_pad:        [B, M, 2, 3] padded GT segments
+        gt_mask:       [B, M] bool mask (True = valid GT segment)
+        eps:           sinkhorn regularization
+        iters:         sinkhorn iterations
+        dustbin_cost:  cost for unmatched segments (scalar or [B])
+        pred_mass:     [B, S] per-segment mass weights (e.g. sigmoid(conf)).
+                       If None, uniform masses are used.
+
+    Returns:
+        [B] per-sample sinkhorn transport cost
+    """
+    B, S, _, _ = pred_segments.shape
+    M = gt_pad.shape[1]
+
+    # Allow per-sample dustbin cost
+    dc = torch.as_tensor(dustbin_cost, device=pred_segments.device, dtype=pred_segments.dtype)
+    if dc.dim() == 0:
+        dc = dc.expand(B)
+
+    # Compute cost matrices [B, S, M] in midpoint-halfvec space.
+    # Decouples position from direction: mid gradient is pure position,
+    # half gradient is pure direction/length. Sign-invariance on half
+    # handles segment direction ambiguity cleanly.
+    p0 = pred_segments[:, :, 0]  # [B, S, 3]
+    p1 = pred_segments[:, :, 1]  # [B, S, 3]
+    g0 = gt_pad[:, :, 0]         # [B, M, 3]
+    g1 = gt_pad[:, :, 1]         # [B, M, 3]
+
+    mid_pred = 0.5 * (p0 + p1)   # [B, S, 3]
+    half_pred = 0.5 * (p1 - p0)  # [B, S, 3]
+    mid_gt = 0.5 * (g0 + g1)     # [B, M, 3]
+    half_gt = 0.5 * (g1 - g0)    # [B, M, 3]
+
+    # Midpoint distance [B, S, M]
+    d_mid = torch.linalg.norm(
+        mid_pred.unsqueeze(2) - mid_gt.unsqueeze(1), dim=-1)
+
+    # Decoupled direction + length distance (sign-invariant for direction ambiguity)
+    len_pred = torch.linalg.norm(half_pred, dim=-1, keepdim=True).clamp(min=1e-6)  # [B, S, 1]
+    len_gt = torch.linalg.norm(half_gt, dim=-1, keepdim=True).clamp(min=1e-6)      # [B, M, 1]
+    dir_pred = half_pred / len_pred  # [B, S, 3]
+    dir_gt = half_gt / len_gt        # [B, M, 3]
+
+    # Direction distance: 1 - |cos(angle)|, sign-invariant [B, S, M]
+    cos_angle = (dir_pred.unsqueeze(2) * dir_gt.unsqueeze(1)).sum(dim=-1)  # [B, S, M]
+    d_dir = 1.0 - cos_angle.abs()
+
+    # Length distance [B, S, M]
+    d_len = (len_pred.unsqueeze(2) - len_gt.unsqueeze(1)).squeeze(-1).abs()
+
+    cost = d_mid + d_dir + d_len  # [B, S, M]
+
+    # Mask invalid GT segments with high cost so they go to dustbin
+    cost = torch.where(gt_mask.unsqueeze(1), cost, dc[:, None, None] * 10.0)
+
+    # Pad with dustbin row and column: [B, S+1, M+1]
+    cost_pad = dc[:, None, None].expand(B, S + 1, M + 1).clone()
+    cost_pad[:, :S, :M] = cost
+    cost_pad[:, -1, -1] = 0.0
+
+    # Masses
+    gt_counts = gt_mask.sum(dim=1).float()  # [B]
+
+    if pred_mass is not None:
+        # Confidence-weighted masses (matches learned_v2 approach).
+        # sigmoid(conf) gives per-segment mass; dustbin masses balance the totals.
+        # No normalization -- sum(a) == sum(b) == max(sum_pred, sum_gt).
+        pm = pred_mass.clamp(min=0.0)  # [B, S]
+        sum_pred = pm.sum(dim=1)       # [B]
+        sum_gt = gt_counts             # [B]
+        pred_dustbin = (sum_gt - sum_pred).clamp(min=0.0)  # [B]
+        gt_dustbin = (sum_pred - sum_gt).clamp(min=0.0)    # [B]
+        a = torch.cat([pm, pred_dustbin.unsqueeze(1)], dim=1)  # [B, S+1]
+        b_val = torch.zeros(B, M + 1, device=cost.device, dtype=cost.dtype)
+        b_val[:, :M] = gt_mask.float()  # 1.0 per valid GT segment
+        b_val[:, -1] = gt_dustbin
+    else:
+        # Uniform masses (normalized)
+        n = float(S)
+        denom = n + gt_counts  # [B]
+        a = (1.0 / denom).unsqueeze(1).expand(B, S + 1).clone()  # [B, S+1]
+        a[:, -1] = gt_counts / denom
+        b_val = (1.0 / denom).unsqueeze(1).expand(B, M + 1).clone()  # [B, M+1]
+        b_val[:, -1] = n / denom
+        # Zero out mass for invalid GT
+        b_val[:, :M] = b_val[:, :M] * gt_mask.float()
+
+    # Log-domain sinkhorn
+    log_a = torch.log(a + 1e-9)
+    log_b = torch.log(b_val + 1e-9)
+    log_k = -cost_pad / eps
+
+    log_u = torch.zeros_like(a)
+    log_v = torch.zeros_like(b_val)
+
+    for _ in range(iters):
+        log_u = log_a - torch.logsumexp(log_k + log_v.unsqueeze(1), dim=2)
+        log_v = log_b - torch.logsumexp(log_k + log_u.unsqueeze(2), dim=1)
+
+    transport = torch.exp(log_u.unsqueeze(2) + log_v.unsqueeze(1) + log_k)
+    return (transport * cost_pad).sum(dim=(1, 2))  # [B]

s23dr_2026_example/tokenizer.py

@@ -0,0 +1,88 @@
+"""Tokenizer: learned embeddings + Fourier features for the point cloud tokens.
+
+The EdgeDepthSequenceBuilder holds the learned embedding tables (label, source,
+behind) and the random Fourier positional encoding. At training time,
+build_tokens() in data.py applies these to pre-sampled point indices on GPU.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .point_fusion import NUM_ADE, NUM_GEST
+
+
+# -- Config --
+
+@dataclass(frozen=True)
+class EdgeDepthSequenceConfig:
+    seq_len: int = 2048
+    colmap_points: int = 1280
+    depth_points: int = 768
+    use_fourier: bool = True
+    fourier_dim: int = 32
+    fourier_scale: float = 10.0
+
+
+# -- Fourier positional encoding --
+
+class FourierFeatures(nn.Module):
+    def __init__(self, in_dim: int = 3, fourier_dim: int = 64,
+                 scale: float = 10.0, seed: int = 0,
+                 learnable: bool = False):
+        super().__init__()
+        gen = torch.Generator()
+        gen.manual_seed(seed)
+        B = torch.randn(fourier_dim, in_dim, generator=gen) * scale
+        if learnable:
+            self.B = nn.Parameter(B)
+        else:
+            self.register_buffer("B", B, persistent=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        proj = (2.0 * np.pi) * (x @ self.B.t())
+        return torch.cat([torch.sin(proj), torch.cos(proj)], dim=-1)
+
+
+# -- Sequence builder (holds embeddings) --
+
+class EdgeDepthSequenceBuilder(nn.Module):
+    """Holds learned embeddings for point cloud tokenization.
+
+    Used by the model at training time: build_tokens() calls
+    self.label_emb(class_id), self.src_emb(source), etc.
+    """
+
+    def __init__(self, cfg: EdgeDepthSequenceConfig, label_emb_dim: int = 16,
+                 src_emb_dim: int = 2, behind_emb_dim: int = 8,
+                 fourier_seed: int = 0, use_vote_features: bool = False,
+                 learnable_fourier: bool = False):
+        super().__init__()
+        self.cfg = cfg
+
+        self.num_labels = 13  # 11 structural + other_house + non_house
+        self.label_emb = nn.Embedding(self.num_labels, label_emb_dim)
+        self.src_emb = nn.Embedding(2, src_emb_dim)
+        self.behind_emb_dim = behind_emb_dim
+        if behind_emb_dim > 0:
+            self.behind_emb = nn.Embedding(NUM_GEST + 1, behind_emb_dim)
+
+        # Fourier positional encoding
+        if cfg.use_fourier:
+            self.pos_enc = FourierFeatures(
+                in_dim=3, fourier_dim=cfg.fourier_dim,
+                scale=cfg.fourier_scale, seed=fourier_seed,
+                learnable=learnable_fourier,
+            )
+            pos_dim = 3 + 2 * cfg.fourier_dim
+        else:
+            self.pos_enc = None
+            pos_dim = 3
+
+        vote_dim = 2 if use_vote_features else 0  # n_views_voted + vote_frac
+        self.use_vote_features = use_vote_features
+        self.out_dim = pos_dim + label_emb_dim + src_emb_dim + behind_emb_dim + vote_dim

s23dr_2026_example/train.py

@@ -0,0 +1,530 @@
+#!/usr/bin/env python3
+"""
+Training script for S23DR 2026.
+
+Usage:
+  python -m s23dr_2026_example.train --cache-dir hf://usm3d/s23dr-2026-sampled_2048_v2:train --steps 80000 --aug-rotate
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path as _Path
+if __package__ is None or __package__ == "":
+    _here = _Path(__file__).resolve().parent
+    if str(_here.parent) not in sys.path:
+        sys.path.insert(0, str(_here.parent))
+    __package__ = _here.name
+
+import argparse
+import gc
+import json
+import math
+import subprocess
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .tokenizer import EdgeDepthSequenceConfig
+from .model import EdgeDepthSegmentsModel
+from .data import build_loader, build_tokens
+from .losses import compute_loss, _loss_inner
+
+# Re-export for eval scripts
+from .data import HFCachedDataset, collate as _collate  # noqa: F401
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    p = argparse.ArgumentParser(description="S23DR 2026 training")
+    p.add_argument("--cache-dir", default=None, help="HF dataset path (hf://repo:split)")
+    p.add_argument("--val-cache-dir", default="", help="Separate cache for validation")
+    p.add_argument("--seq-len", type=int, default=2048,
+                   help="Input sequence length (2048 or 4096, must match dataset)")
+    p.add_argument("--arch", choices=["perceiver", "transformer"], default="perceiver",
+                   help="perceiver=latent bottleneck, transformer=full self-attention encoder")
+    p.add_argument("--segments", type=int, default=32)
+    p.add_argument("--hidden", type=int, default=128)
+    p.add_argument("--ff", type=int, default=512)
+    p.add_argument("--latent-tokens", type=int, default=128)
+    p.add_argument("--latent-layers", type=int, default=7)
+    p.add_argument("--encoder-layers", type=int, default=4,
+                   help="Encoder layers (transformer arch only)")
+    p.add_argument("--pre-encoder-layers", type=int, default=0,
+                   help="Self-attn layers on full token sequence before perceiver bottleneck")
+    p.add_argument("--decoder-layers", type=int, default=3)
+    p.add_argument("--decoder-input-xattn", action="store_true",
+                   help="Add cross-attention from segment queries to input tokens in each decoder layer")
+    p.add_argument("--qk-norm", action="store_true",
+                   help="Normalize Q and K per-head with learned temperature (stabilizes wide models)")
+    p.add_argument("--qk-norm-type", choices=["l2", "rms"], default="l2",
+                   help="QK-norm type: l2 (unit sphere) or rms (RMSNorm, preserves magnitudes)")
+    p.add_argument("--learnable-fourier", action="store_true",
+                   help="Make Fourier positional encoding learnable (vs fixed random)")
+    p.add_argument("--num-heads", type=int, default=4, help="Attention heads")
+    p.add_argument("--kv-heads-cross", type=int, default=2,
+                   help="KV heads for cross-attention (GQA; 0 = standard MHA)")
+    p.add_argument("--kv-heads-self", type=int, default=2,
+                   help="KV heads for self-attention (GQA; 0 = standard MHA)")
+    p.add_argument("--cross-attn-interval", type=int, default=4,
+                   help="Perceiver cross-attention frequency (every N latent layers)")
+    p.add_argument("--dropout", type=float, default=0.1)
+    p.add_argument("--weight-decay", type=float, default=0.01, help="AdamW weight decay")
+    p.add_argument("--steps", type=int, default=5000)
+    p.add_argument("--batch-size", type=int, default=32)
+    p.add_argument("--lr", type=float, default=3e-4)
+    p.add_argument("--adam-betas", default="0.9,0.95", help="AdamW beta1,beta2")
+    p.add_argument("--warmup", type=int, default=200, help="LR warmup steps")
+    p.add_argument("--cosine-decay", action="store_true",
+                   help="Cosine decay LR after warmup (to lr*0.01 at end)")
+    p.add_argument("--cooldown-start", type=int, default=0,
+                   help="Step to begin linear cooldown to lr*0.01 (0=disabled, constant LR after warmup)")
+    p.add_argument("--cooldown-steps", type=int, default=0,
+                   help="Number of steps for linear cooldown (0=no cooldown)")
+    p.add_argument("--seed", type=int, default=7)
+    p.add_argument("--deterministic", action="store_true",
+                   help="Force deterministic mode (disables torch.compile, slower but bit-reproducible)")
+    p.add_argument("--varifold-weight", type=float, default=0.0)
+    p.add_argument("--varifold-cross-only", action="store_true",
+                   help="Drop varifold self-energy (avoids O(S^2) spike, sinkhorn handles repulsion)")
+    p.add_argument("--sinkhorn-weight", type=float, default=1.0)
+    p.add_argument("--sinkhorn-eps", type=float, default=0.1,
+                   help="Sinkhorn regularization (larger = softer matching, stronger gradients)")
+    p.add_argument("--sinkhorn-eps-start", type=float, default=None,
+                   help="Starting eps for epsilon annealing (anneals to --sinkhorn-eps). None=no annealing.")
+    p.add_argument("--sinkhorn-eps-schedule", choices=["linear", "sqrt", "none"], default="none",
+                   help="Eps annealing schedule: linear, sqrt, or none (default: no annealing)")
+    p.add_argument("--sinkhorn-iters", type=int, default=20,
+                   help="Sinkhorn iterations")
+    p.add_argument("--sinkhorn-dustbin", type=float, default=0.3,
+                   help="Sinkhorn dustbin cost in normalized space")
+    p.add_argument("--endpoint-weight", type=float, default=0.0,
+                   help="Weight for endpoint distance loss (sinkhorn-matched, symmetric)")
+    p.add_argument("--endpoint-warmup", type=int, default=0,
+                   help="Steps to linearly warm up endpoint weight from 0 (0=instant)")
+    p.add_argument("--aug-rotate", action="store_true")
+    p.add_argument("--aug-jitter", type=float, default=0.0,
+                   help="Point position jitter std in normalized space (0=disabled, try 0.005)")
+    p.add_argument("--aug-drop", type=float, default=0.0,
+                   help="Fraction of points to randomly drop (0=disabled, try 0.1)")
+    p.add_argument("--aug-flip", action="store_true",
+                   help="Random mirror along X axis (50%% chance)")
+    p.add_argument("--rms-norm", action="store_true", default=True,
+                   help="Use RMSNorm (default). Use --no-rms-norm for LayerNorm")
+    p.add_argument("--no-rms-norm", dest="rms_norm", action="store_false")
+    p.add_argument("--activation", default="gelu", help="FFN activation: gelu, relu, relu_sq")
+    p.add_argument("--behind-emb-dim", type=int, default=8,
+                   help="Behind-gestalt embedding dim (0 to disable)")
+    p.add_argument("--vote-features", action="store_true",
+                   help="Add n_views_voted + vote_frac as raw token features (requires v2 data)")
+    p.add_argument("--segment-param", choices=["midpoint_halfvec", "midpoint_dir_len"],
+                   default="midpoint_halfvec",
+                   help="Output parameterization: halfvec (default) or decoupled direction+length")
+    p.add_argument("--length-floor", type=float, default=0.0,
+                   help="Minimum segment length for midpoint_dir_len (0=no floor)")
+    p.add_argument("--segment-conf", action="store_true",
+                   help="Add per-segment confidence head (use with --conf-thresh at eval)")
+    p.add_argument("--conf-weight", type=float, default=0.0,
+                   help="Weight for confidence loss (requires --segment-conf)")
+    p.add_argument("--conf-mode", choices=["sinkhorn", "sinkhorn_detach"], default="sinkhorn",
+                   help="Confidence training: 'match'=BCE, 'sinkhorn'=OT mass, 'sinkhorn_detach'=OT mass (detached)")
+    p.add_argument("--conf-clamp-min", type=float, default=None,
+                   help="Clamp conf logits to this minimum before sigmoid (e.g., -5)")
+    p.add_argument("--conf-head-wd", type=float, default=None,
+                   help="Separate weight decay for conf head (default: same as other params)")
+    p.add_argument("--ema-decay", type=float, default=0.0,
+                   help="EMA decay rate (0=disabled, try 0.9999). Saves EMA weights in checkpoints.")
+    p.add_argument("--out-dir", default=str(_Path(__file__).resolve().parent / "runs"))
+    p.add_argument("--resume", default="")
+    p.add_argument("--cpu", action="store_true")
+    p.add_argument("--args-from", default=None,
+                   help="Load defaults from a run's args.json (CLI flags override)")
+
+    # If --args-from is specified, load defaults from that JSON file first,
+    # then let CLI flags override.
+    raw_args = p.parse_args()
+    if raw_args.args_from is not None:
+        import json as _json
+        args_path = _Path(raw_args.args_from)
+        if not args_path.exists():
+            raise FileNotFoundError(f"--args-from file not found: {args_path}")
+        saved = _json.loads(args_path.read_text())
+        valid_dests = {a.dest for a in p._actions}
+        defaults = {}
+        for k, v in saved.items():
+            if k in valid_dests and k != "args_from":
+                defaults[k] = v
+        p.set_defaults(**defaults)
+        args = p.parse_args()
+        print(f"Loaded defaults from {args_path} (CLI flags override)")
+    else:
+        args = raw_args
+
+    # Validate required args
+    if not args.cache_dir:
+        p.error("--cache-dir is required (either directly or via --args-from)")
+
+    # Validate arg compatibility
+    if args.arch == "transformer":
+        perceiver_only = []
+        if args.latent_tokens != 128:
+            perceiver_only.append(f"--latent-tokens={args.latent_tokens}")
+        if args.latent_layers != 7:
+            perceiver_only.append(f"--latent-layers={args.latent_layers}")
+        if args.pre_encoder_layers != 0:
+            perceiver_only.append(f"--pre-encoder-layers={args.pre_encoder_layers}")
+        if args.cross_attn_interval != 4:
+            perceiver_only.append(f"--cross-attn-interval={args.cross_attn_interval}")
+        if perceiver_only:
+            raise ValueError(
+                f"Args {', '.join(perceiver_only)} have no effect with --arch transformer. "
+                f"Use --arch perceiver or remove them.")
+    if args.conf_weight > 0 and not args.segment_conf:
+        raise ValueError("--conf-weight requires --segment-conf")
+    if args.conf_mode in ("sinkhorn", "sinkhorn_detach") and args.sinkhorn_weight == 0:
+        raise ValueError("--conf-mode sinkhorn requires --sinkhorn-weight > 0")
+    if args.cosine_decay and args.cooldown_start > 0:
+        raise ValueError("--cosine-decay and --cooldown-start are mutually exclusive")
+
+    device = torch.device("cpu" if args.cpu else ("cuda" if torch.cuda.is_available() else "cpu"))
+    print(f"Device: {device}")
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    # Output
+    import hashlib, os
+    args_hash = hashlib.md5(json.dumps(vars(args), sort_keys=True).encode()).hexdigest()[:4]
+    run_tag = time.strftime("%Y%m%d_%H%M%S") + f"_{args_hash}_{os.getpid() % 10000:04d}"
+    out_dir = Path(args.out_dir) / run_tag
+    out_dir.mkdir(parents=True, exist_ok=True)
+    (out_dir / "checkpoints").mkdir(exist_ok=True)
+
+    # Tee stdout/stderr to run dir
+    import sys as _sys
+    _log_path = out_dir / "train.log"
+    class _Tee:
+        def __init__(self, path, stream):
+            self._file = open(path, "a")
+            self._stream = stream
+        def write(self, data):
+            self._stream.write(data)
+            self._file.write(data)
+            self._file.flush()
+        def flush(self):
+            self._stream.flush()
+            self._file.flush()
+    _sys.stdout = _Tee(_log_path, _sys.stdout)
+    _sys.stderr = _Tee(_log_path, _sys.stderr)
+
+    git_sha = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True,
+                             cwd=str(_Path(__file__).parent)).stdout.strip()
+    git_dirty = subprocess.run(["git", "diff", "--quiet"], capture_output=True,
+                               cwd=str(_Path(__file__).parent)).returncode != 0
+    run_info = {**vars(args), "git_sha": git_sha, "git_dirty": git_dirty}
+    (out_dir / "args.json").write_text(json.dumps(run_info, indent=2, sort_keys=True) + "\n")
+
+    # Set varifold cross-only mode before compile
+    if args.varifold_cross_only:
+        from . import losses as L
+        L.VARIFOLD_CROSS_ONLY = True
+        print("Varifold: cross-only mode (no self-energy)")
+
+    # Model
+    seq_len = args.seq_len
+    norm_class = torch.nn.RMSNorm if args.rms_norm else None
+    seq_cfg = EdgeDepthSequenceConfig(seq_len=seq_len)
+    model = EdgeDepthSegmentsModel(
+        seq_cfg=seq_cfg, segments=args.segments, hidden=args.hidden,
+        num_heads=args.num_heads, kv_heads_cross=args.kv_heads_cross,
+        kv_heads_self=args.kv_heads_self,
+        dim_feedforward=args.ff, dropout=args.dropout,
+        latent_tokens=args.latent_tokens, latent_layers=args.latent_layers,
+        decoder_layers=args.decoder_layers, cross_attn_interval=args.cross_attn_interval,
+        norm_class=norm_class, activation=args.activation,
+        segment_conf=args.segment_conf,
+        segment_param=args.segment_param,
+        length_floor=args.length_floor,
+        arch=args.arch, encoder_layers=args.encoder_layers,
+        pre_encoder_layers=args.pre_encoder_layers,
+        behind_emb_dim=args.behind_emb_dim,
+        use_vote_features=args.vote_features,
+        decoder_input_xattn=args.decoder_input_xattn,
+        qk_norm=args.qk_norm,
+        qk_norm_type=args.qk_norm_type,
+        learnable_fourier=args.learnable_fourier,
+    ).to(device)
+
+    try:
+        from torchinfo import summary
+        summary(model.segmenter,
+                input_data=[torch.zeros(1, seq_len, model.tokenizer.out_dim, device=device),
+                            torch.ones(1, seq_len, device=device, dtype=torch.bool)],
+                col_names=("input_size", "output_size", "num_params"), verbose=1)
+    except ImportError:
+        pass
+    print(f"Total params: {sum(p.numel() for p in model.parameters()):,}")
+
+    # Compile (skip in deterministic mode for bit-reproducibility)
+    torch.set_float32_matmul_precision("high")
+    if args.deterministic:
+        torch.use_deterministic_algorithms(True)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        import os
+        os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":16:8")
+        print("Deterministic mode: no torch.compile, bit-reproducible but ~3x slower")
+    elif device.type == "cuda":
+        model.segmenter = torch.compile(model.segmenter, mode="reduce-overhead", fullgraph=True)
+        from . import losses as L
+        L._loss_fn = torch.compile(_loss_inner, mode="reduce-overhead", fullgraph=True)
+        print("Compiled model + loss (reduce-overhead, fullgraph)")
+
+    # EMA
+    ema_model = None
+    if args.ema_decay > 0:
+        from copy import deepcopy
+        ema_model = deepcopy(model).eval()
+        for p_ema in ema_model.parameters():
+            p_ema.requires_grad_(False)
+        print(f"EMA enabled (decay={args.ema_decay})")
+
+    # Resume
+    start_step = 0
+    if args.resume:
+        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
+        try:
+            model.load_state_dict(ckpt["model"])
+        except RuntimeError:
+            state = {k.replace("segmenter._orig_mod.", "segmenter."): v
+                     for k, v in ckpt["model"].items()}
+            model.load_state_dict(state)
+        start_step = ckpt.get("step", 0)
+        print(f"Resumed from {args.resume} at step {start_step}")
+
+    betas = tuple(float(x) for x in args.adam_betas.split(","))
+
+    # Optimizer: AdamW with optional separate conf_head weight decay
+    conf_wd = args.conf_head_wd if args.conf_head_wd is not None else args.weight_decay
+    if args.conf_head_wd is not None:
+        conf_decay_params = []
+        other_params = []
+        for name, param in model.named_parameters():
+            if not param.requires_grad:
+                continue
+            if 'conf_head' in name:
+                conf_decay_params.append(param)
+            else:
+                other_params.append(param)
+        param_groups = [
+            {"params": other_params, "weight_decay": args.weight_decay},
+            {"params": conf_decay_params, "weight_decay": conf_wd},
+        ]
+        print(f"Conf head WD: {conf_wd} ({len(conf_decay_params)} params)")
+    else:
+        param_groups = model.parameters()
+
+    opt = torch.optim.AdamW(param_groups, lr=args.lr, weight_decay=args.weight_decay,
+                            betas=betas)
+    if args.resume and "optimizer" in ckpt:
+        opt.load_state_dict(ckpt["optimizer"])
+
+    # Data
+    torch.manual_seed(args.seed + 7919)
+    np.random.seed(args.seed + 7919)
+    train_loader = build_loader(args.cache_dir, args.batch_size, aug_rotate=args.aug_rotate,
+                                aug_jitter=args.aug_jitter, aug_drop=args.aug_drop,
+                                aug_flip=args.aug_flip)
+    val_loader = build_loader(args.val_cache_dir, args.batch_size) if args.val_cache_dir else None
+    data_iter = iter(train_loader)
+
+    # Intervals
+    log_int = max(1, min(50, args.steps // 20))
+    ckpt_int = 5000
+    val_int = ckpt_int if val_loader else 0
+
+    # Training loop
+    global_step = start_step
+    loss_ema, loss_sq_ema = 0.0, 0.0
+    t_start = time.perf_counter()
+
+    print(f"Training for {args.steps} steps | {args.segments}seg "
+          f"{args.hidden}h {args.latent_tokens}x{args.latent_layers}L "
+          f"{args.decoder_layers}D")
+
+    # Pre-fetch first batch
+    try:
+        next_batch = next(data_iter)
+    except StopIteration:
+        data_iter = iter(train_loader)
+        next_batch = next(data_iter)
+
+    # Freeze GC after setup to eliminate stalls during training
+    gc.collect()
+    gc.freeze()
+    gc.disable()
+
+    amp_ctx = torch.autocast(device_type='cuda', dtype=torch.bfloat16,
+                             enabled=(device.type == 'cuda'))
+
+    while global_step < args.steps:
+        tokens, masks, gt_list, scales, meta = build_tokens(next_batch, model, device)
+
+        # Epsilon annealing
+        if args.sinkhorn_eps_start is not None and args.sinkhorn_eps_start != args.sinkhorn_eps:
+            if args.sinkhorn_eps_schedule == "sqrt":
+                ratio_sq = (args.sinkhorn_eps_start / args.sinkhorn_eps) ** 2
+                t0 = max(args.steps * 0.8 / max(ratio_sq - 1, 1e-6), 1.0)
+                current_eps = args.sinkhorn_eps_start / math.sqrt(1 + global_step / t0)
+                current_eps = max(current_eps, args.sinkhorn_eps)
+            else:
+                frac = min(global_step / max(args.steps * 0.8, 1), 1.0)
+                current_eps = args.sinkhorn_eps_start + frac * (args.sinkhorn_eps - args.sinkhorn_eps_start)
+        else:
+            current_eps = args.sinkhorn_eps
+
+        with amp_ctx:
+            out = model.forward_tokens(tokens, masks)
+            pred = out["segments"]
+            conf = out.get("conf")
+
+            # Endpoint weight warmup
+            if args.endpoint_warmup > 0 and global_step < args.endpoint_warmup:
+                current_ep_w = args.endpoint_weight * global_step / args.endpoint_warmup
+            else:
+                current_ep_w = args.endpoint_weight
+
+            loss, terms = compute_loss(pred, gt_list, scales.to(device), device,
+                                       args.varifold_weight, args.sinkhorn_weight,
+                                       endpoint_w=current_ep_w,
+                                       conf_logits=conf, conf_weight=args.conf_weight,
+                                       conf_mode=args.conf_mode,
+                                       sinkhorn_eps=current_eps,
+                                       sinkhorn_iters=args.sinkhorn_iters,
+                                       sinkhorn_dustbin=args.sinkhorn_dustbin,
+                                       conf_clamp_min=args.conf_clamp_min)
+
+        loss_val = loss.item()
+        # Adaptive loss spike detection
+        if global_step < 100:
+            loss_ema = loss_val if global_step == start_step else 0.9 * loss_ema + 0.1 * loss_val
+            loss_sq_ema = loss_val**2 if global_step == start_step else 0.9 * loss_sq_ema + 0.1 * loss_val**2
+        else:
+            loss_ema = 0.99 * loss_ema + 0.01 * loss_val
+            loss_sq_ema = 0.99 * loss_sq_ema + 0.01 * loss_val**2
+        loss_std = max(math.sqrt(max(loss_sq_ema - loss_ema**2, 0)), 1e-6)
+        spike_thresh = loss_ema + 5 * loss_std
+
+        # Skip on total loss spike or NaN
+        if not math.isfinite(loss_val) or loss_val > max(spike_thresh, 0.5):
+            sample_ids = [m.get("sample_id", "?") for m in meta]
+            skip_reason = f"loss={loss_val:.2f} > thresh={spike_thresh:.2f}"
+            print(f"Step {global_step}: {skip_reason}, skipping (samples: {sample_ids[:3]})")
+            with open(out_dir / "skipped_samples.jsonl", "a") as f:
+                f.write(json.dumps({"step": global_step, "reason": skip_reason,
+                                    "samples": sample_ids}) + "\n")
+            try:
+                next_batch = next(data_iter)
+            except StopIteration:
+                data_iter = iter(train_loader)
+                next_batch = next(data_iter)
+            continue
+
+        opt.zero_grad()
+        loss.backward()
+
+        # Fetch next batch while GPU finishes backward
+        try:
+            next_batch = next(data_iter)
+        except StopIteration:
+            data_iter = iter(train_loader)
+            next_batch = next(data_iter)
+
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+
+        # LR schedule: warmup -> constant -> optional cooldown or cosine
+        if global_step < args.warmup:
+            lr = args.lr * (global_step + 1) / max(1, args.warmup)
+        elif args.cosine_decay:
+            progress = (global_step - args.warmup) / max(1, args.steps - args.warmup)
+            lr = args.lr * (0.01 + 0.99 * 0.5 * (1 + math.cos(math.pi * progress)))
+        elif args.cooldown_start > 0 and global_step >= args.cooldown_start:
+            progress = (global_step - args.cooldown_start) / max(1, args.cooldown_steps)
+            lr = args.lr * max(0.01, 1.0 - 0.99 * min(1.0, progress))
+        else:
+            lr = args.lr
+        for pg in opt.param_groups:
+            pg["lr"] = lr
+        opt.step()
+        global_step += 1
+
+        # EMA update
+        if ema_model is not None:
+            decay = args.ema_decay
+            with torch.no_grad():
+                for p_ema, p_model in zip(ema_model.parameters(), model.parameters()):
+                    p_ema.lerp_(p_model, 1.0 - decay)
+
+        # Log
+        entry = {"step": global_step, "ts": time.time(), "loss": loss.item(), "lr": lr}
+        entry.update({k: v.item() for k, v in terms.items()})
+        if global_step % log_int == 0:
+            grad_norm = sum(p.grad.norm().item()**2 for p in model.parameters()
+                            if p.grad is not None) ** 0.5
+            entry["grad_norm"] = grad_norm
+
+        if global_step % log_int == 0:
+            ms = (time.perf_counter() - t_start) / log_int * 1000
+            t_start = time.perf_counter()
+            t_str = " ".join(f"{k}={v:.4f}" for k, v in terms.items())
+            print(f"[{global_step}/{args.steps}] loss={loss.item():.4f} {t_str} "
+                  f"lr={lr:.2e} gnorm={entry.get('grad_norm', 0):.3f} [{ms:.0f}ms/step]")
+
+        if val_int > 0 and global_step % val_int == 0:
+            try:
+                vl_list = []
+                with torch.no_grad(), amp_ctx:
+                    for vb in val_loader:
+                        vt, vm, vg, vs, _ = build_tokens(vb, model, device)
+                        vo = model.forward_tokens(vt, vm)
+                        vl, _ = compute_loss(vo["segments"], vg, vs.to(device), device,
+                                             args.varifold_weight, args.sinkhorn_weight)
+                        if math.isfinite(vl.item()):
+                            vl_list.append(vl.item())
+                if vl_list:
+                    val_loss = float(np.mean(vl_list))
+                    print(f"  val_loss={val_loss:.4f}")
+                    entry["val_loss"] = val_loss
+            except Exception as e:
+                print(f"  val eval failed: {e}")
+
+        # Write log entry
+        with open(out_dir / "history.jsonl", "a") as f:
+            f.write(json.dumps(entry) + "\n")
+
+        if global_step % ckpt_int == 0:
+            try:
+                gc.enable(); gc.collect(); gc.freeze(); gc.disable()
+                torch.cuda.empty_cache()
+                save_dict = {"step": global_step, "model": model.state_dict(),
+                             "optimizer": opt.state_dict(), "args": vars(args)}
+                if ema_model is not None:
+                    save_dict["ema_model"] = ema_model.state_dict()
+                torch.save(save_dict, out_dir / "checkpoints" / f"step{global_step:06d}.pt")
+            except Exception as e:
+                print(f"  checkpoint save failed: {e}")
+
+    # Final save
+    save_dict = {"step": global_step, "model": model.state_dict(),
+                 "optimizer": opt.state_dict(), "args": vars(args)}
+    if ema_model is not None:
+        save_dict["ema_model"] = ema_model.state_dict()
+    torch.save(save_dict, out_dir / "checkpoints" / "final.pt")
+    print(f"Done. {global_step} steps. Output: {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

s23dr_2026_example/varifold.py

@@ -0,0 +1,53 @@
+import torch
+
+from .wire_varifold_kernels import (
+    loss_simpson3_batch,
+    loss_simpson3_mix_batch,
+)
+
+
+def segments_to_vertices_edges(segments: torch.Tensor):
+    segs = torch.as_tensor(segments, dtype=torch.float32)
+    vertices = segs.reshape(-1, 3)
+    edges = [(2 * i, 2 * i + 1) for i in range(segs.shape[0])]
+    return vertices, edges
+
+
+def varifold_loss_batch(
+    pred_segments: torch.Tensor,
+    gt_segments: torch.Tensor,
+    *,
+    sigma: float = 0.1,
+    variant: str = "semi_lobatto3",
+    t_nodes01: torch.Tensor | None = None,
+    t_w: torch.Tensor | None = None,
+    sigmas: torch.Tensor | None = None,
+    alpha: torch.Tensor | None = None,
+    normalize_alpha: bool = True,
+    len_pow: float | None = None,
+    gt_mask: torch.Tensor | None = None,
+    pred_weights: torch.Tensor | None = None,
+    cross_only: bool = False,
+) -> torch.Tensor:
+    if pred_segments.dim() != 4 or gt_segments.dim() != 4:
+        raise ValueError("pred_segments and gt_segments must be (B, N, 2, 3)")
+    p_pred, q_pred = pred_segments[:, :, 0], pred_segments[:, :, 1]
+    p_gt, q_gt = gt_segments[:, :, 0], gt_segments[:, :, 1]
+
+    w_gt = None
+    if gt_mask is not None:
+        w_gt = gt_mask.to(device=pred_segments.device, dtype=pred_segments.dtype)
+
+    w_pred = None
+    if pred_weights is not None:
+        w_pred = pred_weights.to(device=pred_segments.device, dtype=pred_segments.dtype)
+
+    if variant != "simpson3":
+        raise ValueError(
+            f"Unsupported varifold variant: {variant!r}. "
+            f"Only 'simpson3' is supported in batch mode.")
+    if sigmas is not None or alpha is not None:
+        if sigmas is None or alpha is None:
+            raise ValueError("sigmas and alpha are required for simpson3 mix")
+        return loss_simpson3_mix_batch(p_pred, q_pred, p_gt, q_gt, sigmas, alpha, w_gt=w_gt, w_pred=w_pred, normalize_alpha=normalize_alpha, cross_only=cross_only)
+    return loss_simpson3_batch(p_pred, q_pred, p_gt, q_gt, sigma, w_gt=w_gt, w_pred=w_pred)

s23dr_2026_example/wire_varifold_kernels.py

@@ -0,0 +1,168 @@
+import torch
+
+# -----------------------------
+# Helpers
+# -----------------------------
+def segment_geom(p: torch.Tensor, q: torch.Tensor, eps: float = 1e-9):
+    """
+    p,q: (...,3)
+    returns d, a, ell, u:
+      d   = q - p
+      a   = ||d||^2
+      ell = sqrt(a + eps^2)
+      u   = d / ell
+    """
+    d = q - p
+    a = (d * d).sum(dim=-1)
+    eps_val = eps
+    if p.dtype in (torch.float16, torch.bfloat16):
+        eps_val = max(eps, float(torch.finfo(p.dtype).eps))
+    ell = torch.sqrt(a + eps_val * eps_val)
+    u = d / ell.unsqueeze(-1)
+    return d, a, ell, u
+
+def sample_points(p: torch.Tensor, q: torch.Tensor, nodes01: torch.Tensor):
+    # (...,3) + (K,) -> (...,K,3)
+    d = q - p
+    nodes = nodes01.to(device=p.device, dtype=p.dtype)
+    shape = [1] * (p.dim() - 1) + [nodes.shape[0], 1]
+    nodes = nodes.view(*shape)
+    return p.unsqueeze(-2) + nodes * d.unsqueeze(-2)
+
+
+# Fixed Lobatto-3 / Simpson nodes+weights on [0,1]
+LOBATTO3_NODES = torch.tensor([0.0, 0.5, 1.0])
+# LOBATTO3_W = torch.tensor([1.0/6.0, 4.0/6.0, 1.0/6.0])
+LOBATTO3_W = torch.tensor([1/3, 1/3, 1/3])
+LOBATTO3_W2 = LOBATTO3_W[:, None] * LOBATTO3_W[None, :]  # (3,3)
+
+
+def _prepare_mix_weights(sigmas, alpha, device, dtype, normalize_alpha: bool):
+    sigmas_t = torch.as_tensor(sigmas, device=device, dtype=dtype).clamp_min(1e-6)
+    alpha_t = torch.as_tensor(alpha, device=device, dtype=dtype)
+    if normalize_alpha:
+        alpha_t = alpha_t / alpha_t.sum().clamp_min(1e-12)
+    return sigmas_t, alpha_t
+
+# -----------------------------
+# Simpson-3 on both segments (3x3 product rule)
+# -----------------------------
+def _prep_weight(w, n: int, b: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor | None:
+    if w is None:
+        return None
+    w = torch.as_tensor(w, device=device, dtype=dtype)
+    if w.dim() == 1:
+        if w.shape[0] != n:
+            raise ValueError(f"weight length {w.shape[0]} != {n}")
+        w = w.unsqueeze(0).expand(b, -1)
+    elif w.dim() == 2:
+        if w.shape[0] != b or w.shape[1] != n:
+            raise ValueError(f"weight shape {tuple(w.shape)} != ({b}, {n})")
+    else:
+        raise ValueError("weights must be 1D or 2D")
+    return w
+
+
+def cross_simpson3(
+    pA,
+    qA,
+    pB,
+    qB,
+    sigma: float | torch.Tensor,
+    wA: torch.Tensor | None = None,
+    wB: torch.Tensor | None = None,
+):
+    device, dtype = pA.device, pA.dtype
+    batched = pA.dim() == 3
+    if not batched:
+        pA = pA.unsqueeze(0)
+        qA = qA.unsqueeze(0)
+        pB = pB.unsqueeze(0)
+        qB = qB.unsqueeze(0)
+    nodes = LOBATTO3_NODES.to(device=device, dtype=dtype)
+    w2 = LOBATTO3_W2.to(device=device, dtype=dtype)
+
+    bsz, nA, _ = pA.shape
+    nB = pB.shape[1]
+    wA = _prep_weight(wA, nA, bsz, device, dtype)
+    wB = _prep_weight(wB, nB, bsz, device, dtype)
+
+    _, _, ellA, uA = segment_geom(pA, qA)
+    _, _, ellB, uB = segment_geom(pB, qB)
+
+    XA = sample_points(pA, qA, nodes)  # (B,N,3,3)
+    YB = sample_points(pB, qB, nodes)  # (B,M,3,3)
+
+    # angular + length factors: (N,M)
+    ang = torch.matmul(uA, uB.transpose(-1, -2)).pow(2)
+    lenfac = ellA[:, :, None] * ellB[:, None, :]
+    if wA is not None or wB is not None:
+        if wA is None:
+            wA = torch.ones((bsz, nA), device=device, dtype=dtype)
+        if wB is None:
+            wB = torch.ones((bsz, nB), device=device, dtype=dtype)
+        lenfac = lenfac * (wA[:, :, None] * wB[:, None, :])
+
+    # spatial: build (N,M,3,3) kernel via broadcasting
+    diff = XA[:, :, None, :, None, :] - YB[:, None, :, None, :, :]  # (B,N,M,3,3,3)
+    r2 = (diff * diff).sum(dim=-1)                                 # (B,N,M,3,3)
+    sigma_t = torch.as_tensor(sigma, device=device, dtype=dtype)
+    if sigma_t.ndim == 0:
+        inv2s2 = 1.0 / (2.0 * sigma_t * sigma_t)
+    else:
+        if sigma_t.shape[0] != bsz:
+            raise ValueError(f"sigma batch {sigma_t.shape[0]} != {bsz}")
+        inv2s2 = (1.0 / (2.0 * sigma_t * sigma_t)).view(bsz, 1, 1, 1, 1)
+    K = torch.exp(-r2 * inv2s2)                                     # (B,N,M,3,3)
+
+    spatial = (K * w2).sum(dim=-1).sum(dim=-1)                     # (B,N,M)
+    out = (ang * lenfac * spatial).sum(dim=-1).sum(dim=-1)         # (B,)
+    return out[0] if not batched else out
+
+
+# -----------------------------
+# Batch losses
+# -----------------------------
+
+def loss_simpson3_batch(
+    p_pred: torch.Tensor,
+    q_pred: torch.Tensor,
+    p_gt: torch.Tensor,
+    q_gt: torch.Tensor,
+    sigma: float | torch.Tensor,
+    w_gt: torch.Tensor | None = None,
+    w_pred: torch.Tensor | None = None,
+    cross_only: bool = False,
+) -> torch.Tensor:
+    cross = cross_simpson3(p_pred, q_pred, p_gt, q_gt, sigma, wA=w_pred, wB=w_gt)
+    if cross_only:
+        # No self-energy: avoids O(S^2) blowup, sinkhorn handles repulsion
+        return -2.0 * cross
+    s_pred = cross_simpson3(p_pred, q_pred, p_pred, q_pred, sigma, wA=w_pred, wB=w_pred)
+    return s_pred - 2.0 * cross
+
+
+def loss_simpson3_mix_batch(
+    p_pred: torch.Tensor,
+    q_pred: torch.Tensor,
+    p_gt: torch.Tensor,
+    q_gt: torch.Tensor,
+    sigmas,
+    alpha,
+    w_gt: torch.Tensor | None = None,
+    w_pred: torch.Tensor | None = None,
+    normalize_alpha: bool = True,
+    cross_only: bool = False,
+) -> torch.Tensor:
+    device, dtype = p_pred.device, p_pred.dtype
+    sigmas_t = torch.as_tensor(sigmas, device=device, dtype=dtype).clamp_min(1e-6)
+    alpha_t = torch.as_tensor(alpha, device=device, dtype=dtype)
+    if normalize_alpha:
+        alpha_t = alpha_t / alpha_t.sum().clamp_min(1e-12)
+    if sigmas_t.ndim == 1:
+        losses = [loss_simpson3_batch(p_pred, q_pred, p_gt, q_gt, s, w_gt=w_gt, w_pred=w_pred, cross_only=cross_only) for s in sigmas_t]
+        return (torch.stack(losses, dim=0) * alpha_t[:, None]).sum(dim=0)
+    if sigmas_t.ndim == 2:
+        losses = [loss_simpson3_batch(p_pred, q_pred, p_gt, q_gt, sigmas_t[:, i], w_gt=w_gt, w_pred=w_pred, cross_only=cross_only) for i in range(sigmas_t.shape[1])]
+        return (torch.stack(losses, dim=0) * alpha_t[:, None]).sum(dim=0)
+    raise ValueError("sigmas must be 1D or 2D for batch loss")

script.py

@@ -0,0 +1,432 @@
+"""S23DR 2026 submission: learned wireframe prediction from fused point clouds.
+
+Pipeline: raw sample -> point fusion -> priority sampling -> model -> post-process -> wireframe
+"""
+from pathlib import Path
+from tqdm import tqdm
+import json
+import os
+import sys
+import time
+
+import numpy as np
+import torch
+
+
+def empty_solution():
+    return np.zeros((2, 3)), [(0, 1)]
+
+
+# ---------------------------------------------------------------------------
+# Point fusion + sampling (from cache_scenes.py / make_sampled_cache.py)
+# ---------------------------------------------------------------------------
+
+# Add our package to path
+SCRIPT_DIR = Path(__file__).resolve().parent
+sys.path.insert(0, str(SCRIPT_DIR))
+
+from s23dr_2026_example.point_fusion import build_compact_scene, FuserConfig
+from s23dr_2026_example.cache_scenes import (
+    _compute_group_and_class, _compute_smart_center_scale,
+)
+from s23dr_2026_example.make_sampled_cache import _priority_sample
+
+# Tokenizer / model imports
+from s23dr_2026_example.tokenizer import EdgeDepthSequenceConfig
+from s23dr_2026_example.model import EdgeDepthSegmentsModel
+from s23dr_2026_example.segment_postprocess import merge_vertices_iterative
+from s23dr_2026_example.varifold import segments_to_vertices_edges
+from s23dr_2026_example.postprocess_v2 import snap_to_point_cloud, snap_horizontal
+from solution import predict_wireframe, read_colmap_rec
+import edge_classifier as ec
+import vertex_refiner as vr
+
+
+def snap_to_handcrafted(learned_v, learned_e, craft_v, snap_radius=1.0, colmap_guard_thresh=2.0):
+    """Move each learned vertex to the nearest handcrafted vertex if within snap_radius.
+
+    Skips entirely if the median nearest-neighbour distance from handcrafted vertices to
+    learned vertices exceeds colmap_guard_thresh -- indicates a COLMAP frame misalignment
+    where the handcrafted pipeline produced vertices in the wrong coordinate frame.
+    """
+    from scipy.spatial.distance import cdist
+
+    if craft_v is None or len(craft_v) == 0 or len(learned_v) == 0:
+        return learned_v, list(learned_e)
+
+    dists = cdist(learned_v, craft_v)  # (N_learned, N_craft)
+
+    median_craft_to_learned = np.median(np.min(dists, axis=0))
+    if median_craft_to_learned > colmap_guard_thresh:
+        print(f"  [COLMAP guard] skipping snap -- median craft->learned dist {median_craft_to_learned:.2f}m")
+        return learned_v, list(learned_e)
+
+    # 1-to-1 greedy matching: sort candidate pairs by distance, each craft vertex used at most once
+    snapped_v = learned_v.copy()
+    used_learned = set()
+    used_craft = set()
+    rows, cols = np.where(dists <= snap_radius)
+    if len(rows) > 0:
+        order = np.argsort(dists[rows, cols])
+        for idx in order:
+            i, j = int(rows[idx]), int(cols[idx])
+            if i not in used_learned and j not in used_craft:
+                # Blend the learned vertex toward the handcrafted one (alpha=0.7).
+                snapped_v[i] = 0.3 * snapped_v[i] + 0.7 * craft_v[j]
+                used_learned.add(i)
+                used_craft.add(j)
+
+    seen = set()
+    new_edges = []
+    for u, v in learned_e:
+        if u == v:
+            continue
+        key = tuple(sorted((u, v)))
+        if key not in seen:
+            seen.add(key)
+            new_edges.append((u, v))
+
+    return snapped_v, new_edges
+
+
+SEQ_LEN = 8192
+COLMAP_QUOTA = 6144
+DEPTH_QUOTA = 2048
+CONF_THRESH = 0.5      # segment-confidence filter threshold
+MERGE_END = 0.6        # iterative vertex-merge end threshold
+SNAP_RADIUS = 0.5      # point-cloud snap radius (metres)
+
+# Classifier-gated handcrafted augmentation. Disabled for the raw model entry;
+# set AUGMENT_ENABLED / VERTEX_AUGMENT_ENABLED to True to enable the hybrid.
+AUGMENT_ENABLED = False
+AUGMENT_THRESHOLD = 0.55
+EDGE_CLASSIFIER_PATH = SCRIPT_DIR / "pnet_class_2026.pth"
+
+VERTEX_AUGMENT_ENABLED = False
+VERTEX_AUGMENT_THRESHOLD = 0.55
+VERTEX_REFINER_PATH = SCRIPT_DIR / "vertex_refiner.pth"
+
+
+def fuse_and_sample(sample, cfg, rng):
+    """Run point fusion + priority sampling on a raw dataset sample.
+
+    Returns a dict with xyz_norm, class_id, source, mask, center, scale, etc.
+    ready for model inference. Returns None if fusion fails.
+    """
+    try:
+        scene = build_compact_scene(sample, cfg, rng)
+    except Exception as e:
+        print(f"  Fusion failed: {e}")
+        return None
+
+    xyz = scene["xyz"]
+    source = scene["source"]
+
+    if len(xyz) < 10:
+        return None
+
+    # Compute group_id and class_id (same as cache_scenes.py)
+    behind_id = scene.get("behind_gest_id", np.full(len(xyz), -1, dtype=np.int16))
+    group_id, class_id = _compute_group_and_class(
+        scene["visible_src"], scene["visible_id"], behind_id, source)
+
+    # Normalize
+    center, scale = _compute_smart_center_scale(xyz, source)
+
+    # Priority sample
+    indices, mask = _priority_sample(source, group_id, SEQ_LEN, COLMAP_QUOTA, DEPTH_QUOTA)
+
+    xyz_norm = (xyz[indices] - center) / scale
+
+    result = {
+        "xyz_norm": xyz_norm.astype(np.float32),
+        "class_id": class_id[indices].astype(np.int64),
+        "source": source[indices].astype(np.int64),
+        "mask": mask,
+        "center": center.astype(np.float32),
+        "scale": np.float32(scale),
+    }
+
+    # Optional fields
+    if "behind_gest_id" in scene:
+        behind = np.clip(scene["behind_gest_id"][indices].astype(np.int16), 0, None)
+        result["behind"] = behind.astype(np.int64)
+    if "n_views_voted" in scene:
+        result["n_views_voted"] = scene["n_views_voted"][indices].astype(np.float32)
+    if "vote_frac" in scene:
+        result["vote_frac"] = scene["vote_frac"][indices].astype(np.float32)
+
+    # Visible src/id for snap post-processing
+    result["visible_src"] = scene["visible_src"][indices].astype(np.int64)
+    result["visible_id"] = scene["visible_id"][indices].astype(np.int64)
+
+    return result
+
+
+def load_model(checkpoint_path, device):
+    """Load model from checkpoint."""
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    args = ckpt.get("args", {})
+
+    norm_class = torch.nn.RMSNorm if args.get("rms_norm") else None
+    seq_cfg = EdgeDepthSequenceConfig(
+        seq_len=SEQ_LEN, colmap_points=COLMAP_QUOTA, depth_points=DEPTH_QUOTA)
+
+    model = EdgeDepthSegmentsModel(
+        seq_cfg=seq_cfg,
+        segments=args.get("segments", 64),
+        hidden=args.get("hidden", 256),
+        num_heads=args.get("num_heads", 4),
+        kv_heads_cross=args.get("kv_heads_cross", 2),
+        kv_heads_self=args.get("kv_heads_self", 2),
+        dim_feedforward=args.get("ff", 1024),
+        dropout=args.get("dropout", 0.1),
+        latent_tokens=args.get("latent_tokens", 256),
+        latent_layers=args.get("latent_layers", 7),
+        decoder_layers=args.get("decoder_layers", 3),
+        cross_attn_interval=args.get("cross_attn_interval", 4),
+        norm_class=norm_class,
+        activation=args.get("activation", "gelu"),
+        segment_conf=args.get("segment_conf", True),
+        behind_emb_dim=args.get("behind_emb_dim", 8),
+        use_vote_features=args.get("vote_features", True),
+        arch=args.get("arch", "perceiver"),
+        encoder_layers=args.get("encoder_layers", 4),
+        pre_encoder_layers=args.get("pre_encoder_layers", 0),
+        segment_param=args.get("segment_param", "midpoint_dir_len"),
+        qk_norm=args.get("qk_norm", True),
+    ).to(device)
+
+    # Handle torch.compile _orig_mod prefix
+    state = ckpt["model"]
+    fixed = {k.replace("segmenter._orig_mod.", "segmenter."): v
+             for k, v in state.items()}
+    model.load_state_dict(fixed, strict=True)
+    model.eval()
+    return model
+
+
+def build_tokens_single(sample_dict, model, device):
+    """Build token tensor for a single sample (no DataLoader)."""
+    xyz = torch.as_tensor(sample_dict["xyz_norm"], dtype=torch.float32).unsqueeze(0).to(device)
+    cid = torch.as_tensor(sample_dict["class_id"], dtype=torch.long).unsqueeze(0).to(device)
+    src = torch.as_tensor(sample_dict["source"], dtype=torch.long).unsqueeze(0).to(device)
+    masks = torch.as_tensor(sample_dict["mask"], dtype=torch.bool).unsqueeze(0).to(device)
+
+    B, T, _ = xyz.shape
+    tok = model.tokenizer
+    fourier = tok.pos_enc(xyz.reshape(-1, 3)).reshape(B, T, -1) \
+        if tok.pos_enc is not None else xyz.new_zeros(B, T, 0)
+    parts = [xyz, fourier, tok.label_emb(cid), tok.src_emb(src.clamp(0, 1))]
+
+    if tok.behind_emb_dim > 0:
+        if "behind" in sample_dict:
+            beh = torch.as_tensor(sample_dict["behind"], dtype=torch.long).unsqueeze(0).to(device)
+        else:
+            beh = xyz.new_zeros(B, T, dtype=torch.long)
+        parts.append(tok.behind_emb(beh))
+
+    if tok.use_vote_features:
+        if "n_views_voted" in sample_dict and "vote_frac" in sample_dict:
+            nv = ((torch.as_tensor(sample_dict["n_views_voted"], dtype=torch.float32).unsqueeze(0).to(device) - 2.7) / 1.0).unsqueeze(-1)
+            vf = ((torch.as_tensor(sample_dict["vote_frac"], dtype=torch.float32).unsqueeze(0).to(device) - 0.5) / 0.25).unsqueeze(-1)
+            parts.extend([nv, vf])
+        else:
+            parts.extend([xyz.new_zeros(B, T, 1), xyz.new_zeros(B, T, 1)])
+
+    tokens = torch.cat(parts, dim=-1)
+    return tokens, masks
+
+
+def predict_sample(sample_dict, model, device, return_stats=False):
+    """Run model inference + post-processing on a fused sample.
+
+    Returns (vertices, edges) in world space.
+    If return_stats=True, also returns a dict of confidence diagnostics.
+    """
+    tokens, masks = build_tokens_single(sample_dict, model, device)
+    scale = float(sample_dict["scale"])
+    center = sample_dict["center"]
+
+    with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.float16,
+                                          enabled=(device.type == 'cuda')):
+        out = model.forward_tokens(tokens, masks)
+
+    segs = out["segments"][0].float().cpu()
+    conf = torch.sigmoid(out["conf"][0].float()).cpu().numpy() if "conf" in out else None
+
+    stats = {
+        'n_segs_total': int(len(segs)),
+        'n_segs_kept': 0,
+        'mean_conf_all': float(conf.mean()) if conf is not None else None,
+        'mean_conf_kept': None,
+        'max_conf': float(conf.max()) if conf is not None else None,
+        'top5_conf_mean': float(np.sort(conf)[-5:].mean()) if conf is not None and len(conf) >= 5 else None,
+    }
+
+    # Confidence filter
+    if conf is not None:
+        keep = conf > CONF_THRESH
+        kept_conf = conf[keep]
+        stats['n_segs_kept'] = int(keep.sum())
+        if keep.sum() > 0:
+            stats['mean_conf_kept'] = float(kept_conf.mean())
+        segs = segs[keep]
+    if len(segs) < 1:
+        if return_stats:
+            return *empty_solution(), stats
+        return empty_solution()
+
+    # To world space
+    segs_world = segs.numpy() * scale + center
+
+    # Vertices + edges from segments
+    pv, pe = segments_to_vertices_edges(torch.tensor(segs_world))
+    pv, pe = pv.numpy(), np.array(pe, dtype=np.int32)
+
+    # Merge
+    pv, pe = merge_vertices_iterative(pv, pe, end=MERGE_END)
+
+    # Snap to point cloud
+    xyz_norm = sample_dict["xyz_norm"]
+    mask = sample_dict["mask"]
+    cid = sample_dict["class_id"]
+    xyz_world = xyz_norm[mask] * scale + center
+    cid_valid = cid[mask]
+    pv = snap_to_point_cloud(pv, xyz_world, cid_valid, snap_radius=SNAP_RADIUS)
+
+    # Horizontal snap
+    pv = snap_horizontal(pv, pe)
+
+    if len(pv) < 2 or len(pe) < 1:
+        if return_stats:
+            return *empty_solution(), stats
+        return empty_solution()
+
+    edges = [(int(a), int(b)) for a, b in pe]
+    if return_stats:
+        return pv, edges, stats
+    return pv, edges
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    t_start = time.time()
+
+    # Load params
+    param_path = Path("params.json")
+    with param_path.open() as f:
+        params = json.load(f)
+    print(f"Competition: {params.get('competition_id', '?')}")
+    print(f"Dataset: {params.get('dataset', '?')}")
+
+    # Load test data
+    data_path = Path("/tmp/data")
+    if not data_path.exists():
+        from huggingface_hub import snapshot_download
+        snapshot_download(
+            repo_id=params["dataset"],
+            local_dir="/tmp/data",
+            repo_type="dataset",
+        )
+
+    from datasets import load_dataset
+    data_files = {
+        "validation": [str(p) for p in data_path.rglob("*public*/**/*.tar")],
+        "test": [str(p) for p in data_path.rglob("*private*/**/*.tar")],
+    }
+    print(f"Data files: {data_files}")
+    dataset = load_dataset(
+        str(data_path / "hoho22k_2026_test_x_anon.py"),
+        data_files=data_files,
+        trust_remote_code=True,
+        writer_batch_size=100,
+    )
+    print(f"Loaded: {dataset}")
+
+    # Load model
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+    checkpoint_path = SCRIPT_DIR / "checkpoint_8192.pt"
+    model = load_model(checkpoint_path, device)
+    print(f"Model loaded: {sum(p.numel() for p in model.parameters()):,} params")
+
+    # Edge classifier (None if disabled or missing)
+    edge_model = None
+    if AUGMENT_ENABLED:
+        if EDGE_CLASSIFIER_PATH.exists():
+            try:
+                edge_model = ec.load_pnet_class(str(EDGE_CLASSIFIER_PATH), device=device)
+                print(f"Edge classifier loaded: "
+                      f"{sum(p.numel() for p in edge_model.parameters()):,} params "
+                      f"(augment threshold {AUGMENT_THRESHOLD})")
+            except Exception as e:
+                print(f"  Edge classifier load failed, augment disabled: {e}")
+        else:
+            print(f"  Edge classifier checkpoint not found at {EDGE_CLASSIFIER_PATH}; augment disabled")
+
+    # Vertex classifier (None if disabled or missing)
+    vertex_model = None
+    if VERTEX_AUGMENT_ENABLED:
+        if VERTEX_REFINER_PATH.exists():
+            try:
+                vertex_model = vr.load_vertex_model(str(VERTEX_REFINER_PATH), device=device)
+                print(f"Vertex classifier loaded: "
+                      f"{sum(p.numel() for p in vertex_model.parameters()):,} params "
+                      f"(augment threshold {VERTEX_AUGMENT_THRESHOLD})")
+            except Exception as e:
+                print(f"  Vertex classifier load failed, augment disabled: {e}")
+        else:
+            print(f"  Vertex classifier checkpoint not found at {VERTEX_REFINER_PATH}; augment disabled")
+
+    # Point fusion config
+    cfg = FuserConfig()
+    rng = np.random.RandomState(2718)
+    aug_rng = np.random.RandomState(31415)
+
+    # Process all samples
+    solution = []
+    total_samples = sum(len(dataset[s]) for s in dataset)
+    processed = 0
+
+    for subset_name in dataset:
+        print(f"\nProcessing {subset_name} ({len(dataset[subset_name])} samples)...")
+
+        for sample in tqdm(dataset[subset_name], desc=subset_name):
+            order_id = sample["order_id"]
+
+            # Fuse + sample (learned pipeline)
+            fused = fuse_and_sample(sample, cfg, rng)
+            if fused is None:
+                pred_v, pred_e = empty_solution()
+            else:
+                try:
+                    pred_v, pred_e = predict_sample(fused, model, device)
+                except Exception as e:
+                    print(f"  Predict failed for {order_id}: {e}")
+                    pred_v, pred_e = empty_solution()
+
+            solution.append({
+                "order_id": order_id,
+                "wf_vertices": pred_v.tolist() if isinstance(pred_v, np.ndarray) else pred_v,
+                "wf_edges": [(int(a), int(b)) for a, b in pred_e],
+            })
+            processed += 1
+
+            if processed % 50 == 0:
+                elapsed = time.time() - t_start
+                rate = elapsed / processed
+                remaining = (total_samples - processed) * rate
+                print(f"  [{processed}/{total_samples}] "
+                      f"{elapsed:.0f}s elapsed, ~{remaining:.0f}s remaining")
+
+    # Save
+    with open("submission.json", "w") as f:
+        json.dump(solution, f)
+
+    elapsed = time.time() - t_start
+    print(f"\nDone. {processed} samples in {elapsed:.0f}s ({elapsed/max(processed,1):.1f}s/sample)")
+    print(f"Saved submission.json ({len(solution)} entries)")

solution.py

@@ -0,0 +1,1210 @@
+import io
+import tempfile
+import zipfile
+from collections import defaultdict
+from typing import Tuple, List, Dict
+import cv2
+import numpy as np
+import pycolmap
+from PIL import Image as PImage
+from scipy.spatial.distance import cdist
+from sklearn.cluster import DBSCAN
+
+from hoho2025.color_mappings import ade20k_color_mapping, gestalt_color_mapping
+
+def empty_solution():
+    '''Return a minimal valid solution, i.e. 2 vertices and 1 edge.'''
+    return np.zeros((2,3)), [(0, 1)]
+    
+    
+def read_colmap_rec(colmap_data):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        with zipfile.ZipFile(io.BytesIO(colmap_data), "r") as zf:
+            zf.extractall(tmpdir)
+        rec = pycolmap.Reconstruction(tmpdir)
+        return rec
+
+def _cam_matrix_from_image(img):
+    """Safely extracts R and t from any pycolmap version."""
+    cfW = img.cam_from_world
+    if callable(cfW):
+        cfW = cfW() 
+    try:
+        R = cfW.rotation.matrix()
+        t = cfW.translation
+    except AttributeError:
+        M = np.array(cfW.matrix())
+        R, t = M[:, :3], M[:, 3]
+    return np.array(R, dtype=np.float64), np.array(t, dtype=np.float64)
+
+def convert_entry_to_human_readable(entry):
+    out = {}
+    for k, v in entry.items():
+        if 'colmap' in k and k != 'pose_only_in_colmap':
+            out['colmap_binary'] = read_colmap_rec(v)
+        elif k in ['wf_vertices', 'wf_edges', 'K', 'R', 't', 'depth']:
+            try:
+                out[k] = np.array(v)
+            except ValueError as e:
+                if "inhomogeneous" in str(e):
+                    out[k] = v  
+                else:
+                    raise e
+        else:
+            out[k] = v
+    out['__key__'] = entry.get('order_id', 'unknown_id')
+    return out
+
+
+def get_house_mask(ade20k_seg):
+    """
+    Get a mask of the house in the ADE20K segmentation map.
+    """
+    house_classes_ade20k = [
+        'wall',
+        'house',
+        'building;edifice',
+        'door;double;door',
+        'windowpane;window',
+    ]
+    np_seg = np.array(ade20k_seg)
+    full_mask = np.zeros(np_seg.shape[:2], dtype=np.uint8)
+    for c in house_classes_ade20k:
+        color = np.array(ade20k_color_mapping[c])
+        mask = cv2.inRange(np_seg, color-0.5, color+0.5)
+        full_mask = np.logical_or(full_mask, mask)
+    return full_mask
+
+
+def point_to_segment_dist(pt, seg_p1, seg_p2):
+    """
+    Computes the Euclidean distance from pt to the line segment p1->p2.
+    pt, seg_p1, seg_p2: (x, y) as np.ndarray
+    """
+    if np.allclose(seg_p1, seg_p2):
+        return np.linalg.norm(pt - seg_p1)
+    seg_vec = seg_p2 - seg_p1
+    pt_vec = pt - seg_p1
+    seg_len2 = seg_vec.dot(seg_vec)
+    t = max(0, min(1, pt_vec.dot(seg_vec)/seg_len2))
+    proj = seg_p1 + t*seg_vec
+    return np.linalg.norm(pt - proj)
+
+
+def project_and_filter_colmap_points(
+    colmap_image: pycolmap.Image,
+    colmap_points3D: Dict[int, 'pycolmap.Point3D'],
+    gest_seg_np: np.ndarray,
+    target_seg_colors: Dict[str, np.ndarray],
+    img_height: int,
+    img_width: int,
+    patch_size: int = 25
+) -> Dict[str, List[np.ndarray]]:
+    """
+    Project COLMAP 3D points to 2D and filter them based on Gestalt segmentation.
+    
+    Returns:
+        Dict mapping class names to lists of 2D points that fall in those segmentation regions.
+    """
+    projected_points_by_class = {}
+
+    for point2D in colmap_image.points2D:
+        if point2D.has_point3D():
+            u, v = point2D.xy[0], point2D.xy[1]
+
+            if 0 <= u < img_width and 0 <= v < img_height:
+                half_patch = patch_size // 2
+
+                v_start = max(0, int(round(v)) - half_patch)
+                v_end = min(img_height, int(round(v)) + half_patch + 1)
+                u_start = max(0, int(round(u)) - half_patch)
+                u_end = min(img_width, int(round(u)) + half_patch + 1)
+
+                seg_color_patch = gest_seg_np[v_start:v_end, u_start:u_end]
+
+                for class_name, target_color in target_seg_colors.items():
+                    patch_matches = np.any(np.all(np.abs(seg_color_patch - target_color) <= 1.0, axis=-1))
+                    if patch_matches:
+                        if class_name not in projected_points_by_class:
+                            projected_points_by_class[class_name] = []
+                        projected_points_by_class[class_name].append(np.array([u, v]))
+
+    return projected_points_by_class
+
+
+def cluster_projected_points_to_vertices(
+    projected_points: List[np.ndarray],
+    eps: float,
+    min_samples: int
+) -> List[np.ndarray]:
+    """
+    Cluster projected 2D points using DBSCAN to find vertex candidates.
+    
+    Returns:
+        List of cluster centroids as vertex locations.
+    """
+    if len(projected_points) < min_samples:
+        return []
+
+    X = np.array(projected_points)
+
+    db = DBSCAN(eps=eps, min_samples=min_samples)
+    labels = db.fit_predict(X)
+
+    vertex_centroids = []
+    unique_labels = set(labels)
+    if -1 in unique_labels:
+        unique_labels.remove(-1)
+
+    for label in unique_labels:
+        class_member_mask = (labels == label)
+        cluster_points = X[class_member_mask]
+        centroid = np.mean(cluster_points, axis=0)
+        vertex_centroids.append(centroid)
+    
+    return vertex_centroids
+
+
+def detect_point_class(
+    gest_seg_np: np.ndarray,
+    class_name: str,
+    gestalt_color_mapping: Dict[str, Tuple[int, int, int]]
+) -> List[Dict[str, any]]:
+    """
+    Detect point-like features (vertices) for a given class in the gestalt segmentation.
+    
+    Args:
+        gest_seg_np: Gestalt segmentation image as numpy array
+        class_name: Name of the class to detect (e.g., 'apex', 'eave_end_point')
+        gestalt_color_mapping: Dictionary mapping class names to RGB colors
+    
+    Returns:
+        List of vertex dictionaries with 'xy' coordinates and 'type'
+    """
+    vertices = []
+    
+    if class_name not in gestalt_color_mapping:
+        return vertices
+    
+    class_color = np.array(gestalt_color_mapping[class_name])
+    class_mask = cv2.inRange(gest_seg_np, class_color-0.5, class_color+0.5)
+    
+    if class_mask.sum() > 0:
+        output = cv2.connectedComponentsWithStats(class_mask, 8, cv2.CV_32S)
+        (numLabels, labels, stats, centroids) = output
+        stats, centroids = stats[1:], centroids[1:]
+        
+        for i in range(numLabels-1):
+            vert = {"xy": centroids[i], "type": class_name}
+            vertices.append(vert)
+    
+    return vertices
+
+def verify_edge_mask(pt1, pt2, semantic_mask, min_overlap=0.4):
+    """
+    Draws a line between two points and verifies that it physically lies on top of 
+    the neural network's semantic mask. Kills lines that cross empty space.
+    """
+    canvas = np.zeros_like(semantic_mask)
+
+    pt1_int = (int(round(pt1[0])), int(round(pt1[1])))
+    pt2_int = (int(round(pt2[0])), int(round(pt2[1])))
+    cv2.line(canvas, pt1_int, pt2_int, 255, 3)
+
+    line_pixels = cv2.countNonZero(canvas)
+    if line_pixels == 0:
+        return False
+
+    overlap = cv2.bitwise_and(canvas, semantic_mask)
+    overlap_pixels = cv2.countNonZero(overlap)
+
+    return (overlap_pixels / line_pixels) >= min_overlap
+
+def get_vertices_and_edges_from_segmentation(
+    gest_seg_np: np.ndarray, 
+    point_class_names: List[str],
+    edge_class_names: List[str],
+    colmap_image: pycolmap.Image = None,
+    colmap_points3D: Dict[int, 'pycolmap.Point3D'] = None,
+    edge_th: float = 25.0,
+    min_3d_points_for_vertex: int = 1,
+    vertex_cluster_eps: float = 5.0,
+    use_colmap_for_vertices: bool = True,
+    patch_size: int = 25
+) -> Tuple[List[dict], List[Tuple[int, int]]]:
+    """
+    Identify apex and eave-end vertices, then detect lines for eave/ridge/rake/valley.
+    Now enhanced with COLMAP 3D point projection and DBSCAN clustering for vertex detection.
+    """
+    if not isinstance(gest_seg_np, np.ndarray):
+        gest_seg_np = np.array(gest_seg_np)
+    
+    vertices = []
+    H, W = gest_seg_np.shape[:2]
+
+    colmap_width = colmap_image.camera.width if colmap_image is not None else None
+    colmap_height = colmap_image.camera.height if colmap_image is not None else None
+
+    if colmap_width != W or colmap_height != H:
+        print(f"Warning: colmap image size {colmap_width}x{colmap_height} does not match gestalt segmentation size {W}x{H}")
+
+    if use_colmap_for_vertices and colmap_image is not None and colmap_points3D is not None:
+        try:
+            target_seg_colors = {}
+            for class_name in point_class_names:
+                if class_name in gestalt_color_mapping:
+                    target_seg_colors[class_name] = np.array(gestalt_color_mapping[class_name])
+            
+            projected_points_by_class = project_and_filter_colmap_points(
+                colmap_image, colmap_points3D, gest_seg_np, target_seg_colors, H, W, patch_size=patch_size
+            )
+
+            for class_name in point_class_names:
+                if class_name in projected_points_by_class:
+                    points_for_class = projected_points_by_class[class_name]
+                    class_centroids = cluster_projected_points_to_vertices(
+                        points_for_class, eps=vertex_cluster_eps, min_samples=min_3d_points_for_vertex
+                    )
+                    
+                    for centroid in class_centroids:
+                        vert = {"xy": centroid, "type": class_name}
+                        vertices.append(vert)
+            
+            print(f"Found {len(vertices)} vertices using COLMAP projection and clustering")
+        except Exception as e:
+            print(f"Error using COLMAP for vertex detection: {e}")
+            vertices = []
+
+    if len(vertices) < 2:
+        print("Using fallback method for vertex detection")
+        vertices = []
+
+        for class_name in point_class_names:
+            point_vertices = detect_point_class(gest_seg_np, class_name, gestalt_color_mapping)
+            vertices.extend(point_vertices)
+
+    structural_pts = []
+    structural_idx_map = []
+    for idx, v in enumerate(vertices):
+        structural_pts.append(v['xy'])
+        structural_idx_map.append(idx)
+    structural_pts = np.array(structural_pts)
+
+    connections = []
+    for edge_class in edge_class_names:
+        if edge_class in ['ridge']:
+            allowed_types = ['apex']
+        elif edge_class in ['eave', 'flashing', 'step_flashing']:
+            allowed_types = ['eave_end_point', 'flashing_end_point']
+        else:
+            allowed_types = ['apex', 'eave_end_point', 'flashing_end_point']
+            
+        allowed_pts = []
+        allowed_idx_map = []
+        for orig_idx, v in enumerate(vertices):
+            if v['type'] in allowed_types:
+                allowed_pts.append(v['xy'])
+                allowed_idx_map.append(orig_idx)
+                
+        allowed_pts = np.array(allowed_pts)
+        if len(allowed_pts) < 2:
+            continue
+
+        edge_color = np.array(gestalt_color_mapping[edge_class])
+        mask_raw = cv2.inRange(gest_seg_np, edge_color-0.5, edge_color+0.5)
+        kernel = np.ones((5, 5), np.uint8)
+        mask = cv2.morphologyEx(mask_raw, cv2.MORPH_CLOSE, kernel)
+        if mask.sum() == 0:
+            continue
+
+        output = cv2.connectedComponentsWithStats(mask, 8, cv2.CV_32S)
+        (numLabels, labels, stats, centroids) = output
+        stats, centroids = stats[1:], centroids[1:]
+        label_indices = range(1, numLabels)
+
+        for lbl in label_indices:
+            mask_i = np.zeros_like(mask)
+            mask_i[labels == lbl] = 255
+
+            lines = cv2.HoughLinesP(mask_i, rho=1, theta=np.pi/180, threshold=15, minLineLength=8, maxLineGap=20)
+            
+            if lines is None:
+                continue
+                
+            for line in lines:
+                x1, y1, x2, y2 = line[0]
+                p1 = np.array([x1, y1], dtype=np.float32)
+                p2 = np.array([x2, y2], dtype=np.float32)
+
+                if len(allowed_pts) < 2:
+                    continue
+
+                dists = np.array([
+                    point_to_segment_dist(allowed_pts[i], p1, p2)
+                    for i in range(len(allowed_pts))
+                ])
+
+                near_mask = (dists <= edge_th)
+                near_indices = np.where(near_mask)[0]
+                if len(near_indices) < 2:
+                    continue
+
+                for i in range(len(near_indices)):
+                    for j in range(i+1, len(near_indices)):
+                        idx_a = near_indices[i]
+                        idx_b = near_indices[j]
+                        
+                        vA = allowed_idx_map[idx_a]
+                        vB = allowed_idx_map[idx_b]
+                        
+                        conn = tuple(sorted((vA, vB)))
+                        if conn not in connections:
+                            is_valid_edge = verify_edge_mask(allowed_pts[idx_a], allowed_pts[idx_b], mask, min_overlap=0.3)
+                            
+                            if is_valid_edge:
+                                connections.append(conn)
+
+    return vertices, connections
+
+
+def get_uv_depth(vertices: List[dict],
+                 depth_fitted: np.ndarray,
+                 sparse_depth: np.ndarray,
+                 search_radius: int = 10) -> Tuple[np.ndarray, np.ndarray]:
+    """For each vertex return its (u, v) and a depth value.
+
+    Uses the nearest valid sparse-depth pixel within search_radius of the vertex;
+    falls back to the dense depth_fitted value when no sparse depth is available.
+    """
+    uv = np.array([vert['xy'] for vert in vertices], dtype=np.float32)
+
+    uv_int = np.round(uv).astype(np.int32)
+    H, W = depth_fitted.shape[:2]
+    uv_int[:, 0] = np.clip(uv_int[:, 0], 0, W - 1)
+    uv_int[:, 1] = np.clip(uv_int[:, 1], 0, H - 1)
+
+    vertex_depth = np.zeros(len(vertices), dtype=np.float32)
+    dense_count = 0
+
+    for i, (x_i, y_i) in enumerate(uv_int):
+        x0 = max(0, x_i - search_radius)
+        x1 = min(W, x_i + search_radius + 1)
+        y0 = max(0, y_i - search_radius)
+        y1 = min(H, y_i + search_radius + 1)
+
+        region = sparse_depth[y0:y1, x0:x1]
+        valid_y, valid_x = np.where(region > 0)
+
+        if valid_y.size > 0:
+            global_x = x0 + valid_x
+            global_y = y0 + valid_y
+            dist_sq = (global_x - x_i)**2 + (global_y - y_i)**2
+            min_idx = np.argmin(dist_sq)
+            vertex_depth[i] = region[valid_y[min_idx], valid_x[min_idx]]
+        else:
+            vertex_depth[i] = depth_fitted[y_i, x_i]
+            dense_count += 1
+    return uv, vertex_depth
+
+
+
+def project_vertices_to_3d(uv: np.ndarray, depth_vert: np.ndarray, col_img: pycolmap.Image, colmap_rec: pycolmap.Reconstruction) -> np.ndarray:
+    xy_local = np.ones((len(uv), 3))
+    
+    try:
+        K = col_img.camera.calibration_matrix()
+    except AttributeError:
+        K = colmap_rec.cameras[col_img.camera_id].calibration_matrix()
+        
+    xy_local[:, 0] = (uv[:, 0] - K[0, 2]) / K[0, 0]
+    xy_local[:, 1] = (uv[:, 1] - K[1, 2]) / K[1, 1]
+    vertices_3d_local = xy_local * depth_vert[...,None]
+    
+    R, t = _cam_matrix_from_image(col_img)
+    world_to_cam = np.eye(4)
+    world_to_cam[:3, :3] = R
+    world_to_cam[:3, 3] = t
+    cam_to_world = np.linalg.inv(world_to_cam)
+    
+    vertices_3d_homogeneous = cv2.convertPointsToHomogeneous(vertices_3d_local)
+    vertices_3d = cv2.transform(vertices_3d_homogeneous, cam_to_world)
+    vertices_3d = cv2.convertPointsFromHomogeneous(vertices_3d).reshape(-1, 3)
+    return vertices_3d
+
+
+def create_3d_wireframe_single_image(vertices: List[dict],
+                                     connections: List[Tuple[int, int]],
+                                     depth: PImage.Image,
+                                     colmap_rec: pycolmap.Reconstruction,
+                                     img_id: str,
+                                     ade_seg: PImage.Image) -> np.ndarray:
+    """Lift one image view's 2D vertices to 3D world coordinates.
+
+    Fits the dense depth to the sparse COLMAP depth, reads a depth per vertex,
+    and back-projects. Returns an empty (0, 3) array if there is no sparse depth.
+    """
+    if (len(vertices) < 2) or (len(connections) < 1):
+        print(f'Warning: create_3d_wireframe_single_image called with insufficient vertices/connections for image {img_id}')
+        return np.empty((0, 3))
+
+    depth_fitted, depth_sparse, found_sparse, col_img = get_fitted_dense_depth(
+        depth, colmap_rec, img_id, ade_seg
+    )
+    if not found_sparse or col_img is None:
+        return np.empty((0, 3))
+
+    uv, depth_vert = get_uv_depth(vertices, depth_fitted, depth_sparse, search_radius=25)
+    vertices_3d = project_vertices_to_3d(uv, depth_vert, col_img, colmap_rec)
+    return vertices_3d
+
+
+def merge_vertices_3d(vert_edge_per_image, point_class_names: List[str], th=0.5):
+    '''Merge vertices that are close in 3D space and of the same type.'''
+    all_3d_vertices = []
+    connections_3d = []
+    all_indexes = []
+    cur_start = 0
+    types = []
+
+    type_to_id = {class_name: idx for idx, class_name in enumerate(point_class_names)}
+
+    for cimg_idx, (vertices, connections, vertices_3d) in vert_edge_per_image.items():
+        vertex_type_ids = []
+        for v in vertices:
+            vertex_type = v['type']
+            type_id = type_to_id.get(vertex_type, -1)
+            vertex_type_ids.append(type_id)
+
+        types += vertex_type_ids
+        all_3d_vertices.append(vertices_3d)
+        connections_3d+=[(x+cur_start,y+cur_start) for (x,y) in connections]
+        cur_start+=len(vertices_3d)
+    all_3d_vertices = np.concatenate(all_3d_vertices, axis=0)
+
+    distmat = cdist(all_3d_vertices, all_3d_vertices)
+    types = np.array(types).reshape(-1,1)
+    same_types = cdist(types, types)
+
+    # Merge vertices that are both close in space and of the same type.
+    mask_to_merge = (distmat <= th) & (same_types==0)
+    new_vertices = []
+    new_connections = []
+
+    to_merge = sorted(list(set([tuple(a.nonzero()[0].tolist()) for a in mask_to_merge])))
+
+    # Transitive grouping: union overlapping merge-sets into connected groups.
+    to_merge_final = defaultdict(list)
+    for i in range(len(all_3d_vertices)):
+        for j in to_merge:
+            if i in j:
+                to_merge_final[i]+=j
+
+    for k, v in to_merge_final.items():
+        to_merge_final[k] = list(set(v))
+
+    already_there = set()
+    merged = []
+    for k, v in to_merge_final.items():
+        if k in already_there:
+            continue
+        merged.append(v)
+        for vv in v:
+            already_there.add(vv)
+
+    old_idx_to_new = {}
+    count=0
+    for idxs in merged:
+        new_vertices.append(all_3d_vertices[idxs].mean(axis=0))
+        for idx in idxs:
+            old_idx_to_new[idx] = count
+        count +=1
+    new_vertices=np.array(new_vertices)
+
+    for conn in connections_3d:
+        new_con = sorted((old_idx_to_new[conn[0]], old_idx_to_new[conn[1]]))
+        if new_con[0] == new_con[1]:
+            continue
+        if new_con not in new_connections:
+            new_connections.append(new_con)
+    return new_vertices, new_connections
+
+
+def prune_not_connected(all_3d_vertices, connections_3d, keep_largest=True):
+    """
+    Prune vertices not connected to anything. If keep_largest=True, also
+    keep only the largest connected component in the graph.
+    """
+    if len(all_3d_vertices) == 0:
+        return np.empty((0, 3)), []
+
+    adj = defaultdict(set)
+    for (i, j) in connections_3d:
+        adj[i].add(j)
+        adj[j].add(i)
+
+    used_idxs = set()
+    for (i, j) in connections_3d:
+        used_idxs.add(i)
+        used_idxs.add(j)
+
+    if not used_idxs:
+        return np.empty((0,3)), []
+
+    # If we only want to remove truly isolated points, but keep multiple subgraphs:
+    if not keep_largest:
+        new_map = {}
+        used_list = sorted(list(used_idxs))
+        for new_id, old_id in enumerate(used_list):
+            new_map[old_id] = new_id
+        new_vertices = np.array([all_3d_vertices[old_id] for old_id in used_list])
+        new_conns = []
+        for (i, j) in connections_3d:
+            if i in used_idxs and j in used_idxs:
+                new_conns.append((new_map[i], new_map[j]))
+        return new_vertices, new_conns
+
+    # Otherwise find the largest connected component:
+    visited = set()
+    def bfs(start):
+        queue = [start]
+        comp = []
+        visited.add(start)
+        while queue:
+            cur = queue.pop()
+            comp.append(cur)
+            for neigh in adj[cur]:
+                if neigh not in visited:
+                    visited.add(neigh)
+                    queue.append(neigh)
+        return comp
+
+    comps = []
+    for idx in used_idxs:
+        if idx not in visited:
+            c = bfs(idx)
+            comps.append(c)
+
+    comps.sort(key=lambda c: len(c), reverse=True)
+    largest = comps[0] if len(comps)>0 else []
+
+    new_map = {}
+    for new_id, old_id in enumerate(largest):
+        new_map[old_id] = new_id
+
+    new_vertices = np.array([all_3d_vertices[old_id] for old_id in largest])
+    new_conns = []
+    for (i, j) in connections_3d:
+        if i in largest and j in largest:
+            new_conns.append((new_map[i], new_map[j]))
+
+    new_conns = list(set([tuple(sorted(c)) for c in new_conns]))
+    return new_vertices, new_conns
+
+def get_sparse_depth(colmap_rec, img_id_substring, depth):
+    H, W = depth.shape
+    found_img = None
+    for img_id_c, col_img in colmap_rec.images.items():
+        if img_id_substring in col_img.name:
+            found_img = col_img
+            break
+    if found_img is None:
+        return np.zeros((H, W), dtype=np.float32), False, None
+
+    points_xyz = []
+    for pid, p3D in colmap_rec.points3D.items():
+        if found_img.has_point3D(pid):
+            points_xyz.append(p3D.xyz) 
+    if not points_xyz:
+        return np.zeros((H, W), dtype=np.float32), False, found_img
+
+    points_xyz = np.array(points_xyz) 
+    uv = []
+    z_vals = []
+
+    cam = colmap_rec.cameras[found_img.camera_id]
+    R, t = _cam_matrix_from_image(found_img)
+    K = cam.calibration_matrix()
+
+    for xyz in points_xyz:
+        p_cam = R @ np.asarray(xyz, dtype=np.float64) + t
+        if p_cam[2] > 0:
+            u = p_cam[0] / p_cam[2] * K[0, 0] + K[0, 2]
+            v = p_cam[1] / p_cam[2] * K[1, 1] + K[1, 2]
+            u_i, v_i = int(round(u)), int(round(v))
+            if 0 <= u_i < W and 0 <= v_i < H:
+                uv.append((u_i, v_i))
+                z_vals.append(p_cam[2]) 
+
+    uv = np.array(uv, dtype=int)     
+    z_vals = np.array(z_vals)        
+    
+    depth_out = np.zeros((H, W), dtype=np.float32)
+    if len(uv) > 0:
+        depth_out[uv[:,1], uv[:,0]] = z_vals 
+    
+    return depth_out, True, found_img
+
+
+def fit_scale_robust_median(depth, sparse_depth, validity_mask=None):
+    """
+    Fit a scale factor to the depth map using the median of the ratio of sparse to dense depth.
+    """
+    if validity_mask is None:
+        mask = (sparse_depth != 0)
+    else:
+        mask = (sparse_depth != 0) & validity_mask
+    mask = mask & (depth <50) & (sparse_depth <50)
+    X = depth[mask]
+    Y = sparse_depth[mask]
+    alpha =np.median(Y/X)
+    depth_fitted = alpha * depth
+    return alpha, depth_fitted
+    
+
+def get_fitted_dense_depth(depth, colmap_rec, img_id, ade20k_seg):
+    """Scale the dense depth to align with the sparse COLMAP depth.
+
+    Reads sparse depth from COLMAP, fits a scale factor using only points inside
+    the ADE20k house mask, and returns the scaled dense depth and the sparse
+    depth. found_sparse is False when no sparse depth is available for the image.
+    """
+    depth_np = np.array(depth) / 1000.  # mm to meters
+    depth_sparse, found_sparse, col_img = get_sparse_depth(colmap_rec, img_id, depth_np)
+
+    if not found_sparse:
+        print(f'No sparse depth found for image {img_id}')
+        return depth_np, np.zeros_like(depth_np), False, None
+
+    house_mask = get_house_mask(ade20k_seg)
+    k, depth_fitted = fit_scale_robust_median(depth_np, depth_sparse, validity_mask=house_mask)
+    print(f"Fitted depth scale k={k:.4f} for image {img_id}")
+    return depth_fitted, depth_sparse, True, col_img
+
+
+def precompute_overlapping_views(
+    colmap_reconstruction: pycolmap.Reconstruction,
+    min_shared_points: int = 10
+) -> Dict[int, List[pycolmap.Image]]:
+    """Map each image to the images it shares at least min_shared_points 3D points with.
+
+    Returns a dict image_id -> list of overlapping images (excluding self).
+    """
+    print("Pre-computing overlapping views...")
+
+    image_3d_points = {}
+    for img_id, image in colmap_reconstruction.images.items():
+        points_3d = set()
+        for point2D in image.points2D:
+            if point2D.has_point3D():
+                points_3d.add(point2D.point3D_id)
+        image_3d_points[img_id] = points_3d
+
+    overlapping_views = {}
+    total_pairs = 0
+    overlapping_pairs = 0
+    
+    for img_id_1, image_1 in colmap_reconstruction.images.items():
+        overlapping_views[img_id_1] = []
+        points_1 = image_3d_points[img_id_1]
+        
+        if len(points_1) == 0:
+            continue
+        
+        for img_id_2, image_2 in colmap_reconstruction.images.items():
+            if img_id_1 >= img_id_2:
+                continue
+
+            total_pairs += 1
+            points_2 = image_3d_points[img_id_2]
+
+            shared_points = points_1.intersection(points_2)
+            if len(shared_points) >= min_shared_points:
+                overlapping_pairs += 1
+                overlapping_views[img_id_1].append(image_2)
+                if img_id_2 not in overlapping_views:
+                    overlapping_views[img_id_2] = []
+                overlapping_views[img_id_2].append(image_1)
+    
+    print(f"  Found {overlapping_pairs}/{total_pairs} overlapping pairs")
+    avg_overlaps = np.mean([len(overlaps) for overlaps in overlapping_views.values()]) if overlapping_views else 0
+    print(f"  Average overlaps per image: {avg_overlaps:.1f}")
+    
+    return overlapping_views
+
+
+def check_3d_point_multi_view_consistency(
+    point_3d: np.ndarray,
+    original_vertex_type: str,
+    current_colmap_image: pycolmap.Image,
+    precomputed_overlaps: Dict[int, List[pycolmap.Image]],
+    image_data_map: Dict[str, np.ndarray],
+    gestalt_color_mapping: Dict[str, tuple],
+    min_consistent_views: int = 2,
+    projection_patch_size: int = 5,
+    debug: bool = False
+) -> bool:
+    """Check whether a 3D vertex is consistent across the views that see it.
+
+    Projects the point into each overlapping view and checks that the gestalt
+    segmentation around the projection matches the vertex's class. Returns True
+    if at least min_consistent_views agree, and also True when there are too few
+    overlapping views to verify.
+    """
+    if original_vertex_type not in gestalt_color_mapping:
+        if debug:
+            print(f"  Vertex type {original_vertex_type} not in color mapping")
+        return False
+
+    target_color = np.array(gestalt_color_mapping[original_vertex_type])
+
+    overlapping_views = precomputed_overlaps.get(current_colmap_image.image_id, [])
+
+    if debug:
+        print(f"  Found {len(overlapping_views)} overlapping views for vertex type {original_vertex_type}")
+
+    if len(overlapping_views) < min_consistent_views:
+        if debug:
+            print(f"  Not enough overlapping views ({len(overlapping_views)} < {min_consistent_views}), accepting point")
+        return True  # Accept when we cannot verify (too few overlapping views).
+    
+    consistent_view_count = 0
+    total_checked_views = 0
+    half_patch = projection_patch_size // 2
+    
+    for other_image in overlapping_views:
+        try:
+            total_checked_views += 1
+            projection = other_image.project_point(point_3d)
+            if projection is None:
+                if debug:
+                    print(f"    View {other_image.name}: projection failed (behind camera)")
+                continue
+
+            u, v = projection
+
+            img_width = other_image.camera.width
+            img_height = other_image.camera.height
+            if not (0 <= u < img_width and 0 <= v < img_height):
+                if debug:
+                    print(f"    View {other_image.name}: projection out of bounds ({u:.1f}, {v:.1f})")
+                continue
+
+            other_gest_seg_np = None
+            for img_name, gest_seg_np in image_data_map.items():
+                if img_name in other_image.name or other_image.name in img_name:
+                    other_gest_seg_np = gest_seg_np
+                    break
+
+            if other_gest_seg_np is None:
+                if debug:
+                    print(f"    View {other_image.name}: no segmentation data found")
+                continue
+
+            seg_h, seg_w = other_gest_seg_np.shape[:2]
+
+            # Rescale the projection if the segmentation resolution differs from the camera.
+            if seg_w != img_width or seg_h != img_height:
+                u_seg = int(round(u * seg_w / img_width))
+                v_seg = int(round(v * seg_h / img_height))
+            else:
+                u_seg, v_seg = int(round(u)), int(round(v))
+
+            v_start = max(0, v_seg - half_patch)
+            v_end = min(seg_h, v_seg + half_patch + 1)
+            u_start = max(0, u_seg - half_patch)
+            u_end = min(seg_w, u_seg + half_patch + 1)
+
+            seg_color_patch = other_gest_seg_np[v_start:v_end, u_start:u_end]
+
+            if seg_color_patch.size > 0:
+                patch_matches = np.any(np.all(np.abs(seg_color_patch - target_color) <= 1.0, axis=-1))
+                if patch_matches:
+                    consistent_view_count += 1
+                    if debug:
+                        print(f"    View {other_image.name}: MATCH at ({u:.1f}, {v:.1f}) -> ({u_seg}, {v_seg})")
+                else:
+                    if debug:
+                        unique_colors = np.unique(seg_color_patch.reshape(-1, 3), axis=0)
+                        print(f"    View {other_image.name}: no match at ({u:.1f}, {v:.1f}) -> ({u_seg}, {v_seg}), colors: {unique_colors[:3]}")
+        
+        except Exception as e:
+            if debug:
+                print(f"    View {other_image.name}: exception {e}")
+            continue
+    
+    result = consistent_view_count >= min_consistent_views
+    if debug:
+        print(f"  Result: {consistent_view_count}/{total_checked_views} consistent views, required: {min_consistent_views}, accepted: {result}")
+    
+    return result
+
+
+def filter_vertices_by_multi_view_consistency(
+    vert_edge_per_image: Dict[int, Tuple[List[dict], List[Tuple[int, int]], np.ndarray]],
+    colmap_reconstruction: pycolmap.Reconstruction,
+    gestalt_segmentations: List[PImage.Image],
+    image_ids: List[str],
+    gestalt_color_mapping: Dict[str, tuple],
+    depth_size_per_image: List[Tuple[int, int]],  # [(W, H), ...] for each image
+    min_consistent_views: int = 2,
+    min_shared_points_for_overlap: int = 10,
+    projection_patch_size: int = 25
+) -> Dict[int, Tuple[List[dict], List[Tuple[int, int]], np.ndarray]]:
+    """Drop 3D vertices that are not multi-view consistent and remap edges to the survivors."""
+    precomputed_overlaps = precompute_overlapping_views(
+        colmap_reconstruction, min_shared_points_for_overlap
+    )
+
+    image_data_map = {}
+    for i, (gest_seg, img_id, (w, h)) in enumerate(zip(gestalt_segmentations, image_ids, depth_size_per_image)):
+        gest_seg_resized = gest_seg.resize((w, h))
+        gest_seg_np = np.array(gest_seg_resized).astype(np.uint8)
+        image_data_map[img_id] = gest_seg_np
+
+    filtered_vert_edge_per_image = {}
+
+    for img_idx, (orig_2d_verts, orig_2d_conns, v3d_candidates) in vert_edge_per_image.items():
+        if len(v3d_candidates) == 0:
+            filtered_vert_edge_per_image[img_idx] = ([], [], np.empty((0, 3)))
+            continue
+
+        current_img_id = image_ids[img_idx]
+        current_colmap_img = None
+        for colmap_img_id, colmap_img in colmap_reconstruction.images.items():
+            if current_img_id in colmap_img.name:
+                current_colmap_img = colmap_img
+                break
+
+        if current_colmap_img is None:
+            filtered_vert_edge_per_image[img_idx] = (orig_2d_verts, orig_2d_conns, v3d_candidates)
+            continue
+
+        kept_v3d = []
+        kept_orig_2d_verts_indices = []
+
+        for j, p_3d in enumerate(v3d_candidates):
+            if j >= len(orig_2d_verts):
+                continue
+
+            original_vertex_type = orig_2d_verts[j]['type']
+
+            is_consistent = check_3d_point_multi_view_consistency(
+                p_3d, original_vertex_type, current_colmap_img, precomputed_overlaps,
+                image_data_map, gestalt_color_mapping, min_consistent_views,
+                projection_patch_size=projection_patch_size,
+                debug=False
+            )
+
+            if is_consistent:
+                kept_v3d.append(p_3d)
+                kept_orig_2d_verts_indices.append(j)
+
+        if len(kept_v3d) == 0:
+            filtered_vert_edge_per_image[img_idx] = ([], [], np.empty((0, 3)))
+            continue
+
+        new_orig_2d_verts = [orig_2d_verts[j] for j in kept_orig_2d_verts_indices]
+
+        old_idx_to_new_idx = {old_idx: new_idx for new_idx, old_idx in enumerate(kept_orig_2d_verts_indices)}
+        new_orig_2d_conns = []
+        
+        for (u, v) in orig_2d_conns:
+            if u in old_idx_to_new_idx and v in old_idx_to_new_idx:
+                new_u = old_idx_to_new_idx[u]
+                new_v = old_idx_to_new_idx[v]
+                new_orig_2d_conns.append((new_u, new_v))
+        
+        filtered_vert_edge_per_image[img_idx] = (
+            new_orig_2d_verts, 
+            new_orig_2d_conns, 
+            np.array(kept_v3d)
+        )
+    
+    return filtered_vert_edge_per_image
+
+
+def recover_edges_after_vertex_filtering(
+    filtered_vert_edge_per_image: Dict[int, Tuple[List[dict], List[Tuple[int, int]], np.ndarray]],
+    good_entry: dict,
+    edge_class_names: List[str],
+    edge_th: float = 25.0
+) -> Dict[int, Tuple[List[dict], List[Tuple[int, int]], np.ndarray]]:
+    """Re-detect edges between the surviving vertices after vertex filtering, using the
+    same semantic rulebook and line-of-sight verification as the initial edge detection."""
+    recovered_vert_edge_per_image = {}
+    total_new_edges = 0
+    total_original_edges = 0
+
+    for img_idx, (filtered_2d_verts, filtered_2d_conns, filtered_v3d) in filtered_vert_edge_per_image.items():
+        total_original_edges += len(filtered_2d_conns)
+
+        if len(filtered_2d_verts) < 2:
+            recovered_vert_edge_per_image[img_idx] = (filtered_2d_verts, filtered_2d_conns, filtered_v3d)
+            continue
+
+        try:
+            gest = good_entry['gestalt'][img_idx]
+            depth = good_entry['depth'][img_idx]
+            depth_size = (np.array(depth).shape[1], np.array(depth).shape[0])  # W, H
+            gest_seg = gest.resize(depth_size)
+            gest_seg_np = np.array(gest_seg).astype(np.uint8)
+        except (IndexError, KeyError):
+            recovered_vert_edge_per_image[img_idx] = (filtered_2d_verts, filtered_2d_conns, filtered_v3d)
+            continue
+
+        structural_pts = np.array([v['xy'] for v in filtered_2d_verts])
+        structural_idx_map = list(range(len(filtered_2d_verts)))
+
+        new_connections = []
+        for edge_class in edge_class_names:
+            
+            # --- 1. THE SEMANTIC RULEBOOK ---
+            if edge_class in ['ridge']:
+                allowed_types = ['apex']
+            elif edge_class in ['eave', 'flashing', 'step_flashing']:
+                allowed_types = ['eave_end_point', 'flashing_end_point']
+            else: # rake, valley, hip, transition_line
+                allowed_types = ['apex', 'eave_end_point', 'flashing_end_point']
+                
+            allowed_pts = []
+            allowed_idx_map = []
+            for orig_idx, v in enumerate(filtered_2d_verts):
+                if v['type'] in allowed_types:
+                    allowed_pts.append(v['xy'])
+                    allowed_idx_map.append(orig_idx)
+                    
+            allowed_pts = np.array(allowed_pts)
+            if len(allowed_pts) < 2:
+                continue
+            # --------------------------------
+                
+            edge_color = np.array(gestalt_color_mapping[edge_class])
+            mask_raw = cv2.inRange(gest_seg_np, edge_color-0.5, edge_color+0.5)
+            kernel = np.ones((5, 5), np.uint8)
+            mask = cv2.morphologyEx(mask_raw, cv2.MORPH_CLOSE, kernel)
+            if mask.sum() == 0:
+                continue
+
+            output = cv2.connectedComponentsWithStats(mask, 8, cv2.CV_32S)
+            (numLabels, labels, stats, centroids) = output
+            stats, centroids = stats[1:], centroids[1:]
+            label_indices = range(1, numLabels)
+
+            for lbl in label_indices:
+                mask_i = np.zeros_like(mask)
+                mask_i[labels == lbl] = 255
+                
+                # HoughLinesP for discrete line segments
+                lines = cv2.HoughLinesP(mask_i, rho=1, theta=np.pi/180, threshold=15, minLineLength=8, maxLineGap=20)
+                
+                if lines is None:
+                    continue
+                    
+                for line in lines:
+                    x1, y1, x2, y2 = line[0]
+                    p1 = np.array([x1, y1], dtype=np.float32)
+                    p2 = np.array([x2, y2], dtype=np.float32)
+
+                    if len(allowed_pts) < 2:
+                        continue
+
+                    # Distance check using ONLY the semantically allowed points
+                    dists = np.array([
+                        point_to_segment_dist(allowed_pts[i], p1, p2)
+                        for i in range(len(allowed_pts))
+                    ])
+
+                    near_mask = (dists <= edge_th)
+                    near_indices = np.where(near_mask)[0]
+                    if len(near_indices) < 2:
+                        continue
+
+                    # --- 2. CONNECTIVITY WITH LINE-OF-SIGHT VERIFICATION ---
+                    for i in range(len(near_indices)):
+                        for j in range(i+1, len(near_indices)):
+                            idx_a = near_indices[i]
+                            idx_b = near_indices[j]
+                            
+                            vA = allowed_idx_map[idx_a]
+                            vB = allowed_idx_map[idx_b]
+                            
+                            conn = tuple(sorted((vA, vB)))
+                            if conn not in new_connections:
+                                # THE ULTIMATE SPIDERWEB KILLER:
+                                # Verify that the line between these two corners actually exists in the mask!
+                                is_valid_edge = verify_edge_mask(allowed_pts[idx_a], allowed_pts[idx_b], mask, min_overlap=0.3)
+                                
+                                if is_valid_edge:
+                                    new_connections.append(conn)
+        
+        total_new_edges += len(new_connections)
+        recovered_vert_edge_per_image[img_idx] = (filtered_2d_verts, new_connections, filtered_v3d)
+    
+    print(f"  Edge recovery details: {total_original_edges} -> {total_new_edges} edges across all images")
+    return recovered_vert_edge_per_image
+
+def merge_collinear_edges(vertices: np.ndarray, edges: List[Tuple[int, int]], cos_threshold: float = -0.98) -> Tuple[np.ndarray, List[Tuple[int, int]]]:
+    """
+    Finds degree-2 vertices that form a straight line and merges their edges.
+    cos_threshold of -0.98 corresponds to ~170 degrees.
+    """
+    if len(edges) == 0:
+        return vertices, edges
+        
+    adj = defaultdict(set)
+    for u, v in edges:
+        adj[u].add(v)
+        adj[v].add(u)
+        
+    edges_set = set([tuple(sorted((u, v))) for u, v in edges])
+    
+    merged_something = True
+    while merged_something:
+        merged_something = False
+        
+        for b in list(adj.keys()):
+            neighbors = list(adj[b])
+            if len(neighbors) == 2:
+                a, c = neighbors
+                
+                vec1 = vertices[a] - vertices[b]
+                vec2 = vertices[c] - vertices[b]
+                
+                norm1 = np.linalg.norm(vec1)
+                norm2 = np.linalg.norm(vec2)
+                
+                if norm1 > 1e-5 and norm2 > 1e-5:
+                    cos_sim = np.dot(vec1, vec2) / (norm1 * norm2)
+                    
+                    if cos_sim < cos_threshold:
+                        e1 = tuple(sorted((a, b)))
+                        e2 = tuple(sorted((b, c)))
+                        
+                        if e1 in edges_set: edges_set.remove(e1)
+                        if e2 in edges_set: edges_set.remove(e2)
+                        
+                        new_edge = tuple(sorted((a, c)))
+                        edges_set.add(new_edge)
+                        
+                        adj[a].remove(b)
+                        adj[c].remove(b)
+                        adj[a].add(c)
+                        adj[c].add(a)
+                        del adj[b]
+                        
+                        merged_something = True
+                        break 
+                        
+    new_edges = list(edges_set)
+    return vertices, new_edges
+
+
+def predict_wireframe(entry) -> Tuple[np.ndarray, List[int]]:
+    """Predict the 3D wireframe (vertices and edges) from a dataset entry."""
+    good_entry = convert_entry_to_human_readable(entry)
+    vert_edge_per_image = {}
+
+    depth_sizes = []
+    colmap_rec = good_entry.get('colmap', good_entry.get('colmap_binary'))
+
+    for i, (gest, depth, K, R, t, img_id, ade_seg) in enumerate(zip(good_entry['gestalt'],
+                                                good_entry['depth'],
+                                                good_entry['K'],
+                                                good_entry['R'],
+                                                good_entry['t'],
+                                                good_entry['image_ids'],
+                                                good_entry['ade']
+                                                )):
+
+        K = np.array(K)
+        R = np.array(R)
+        t = np.array(t)
+        depth_size = (np.array(depth).shape[1], np.array(depth).shape[0]) # W, H
+        depth_sizes.append(depth_size)
+
+        # resize() can place pixels at the wrong positions; see
+        # https://numpy.org/doc/stable/reference/generated/numpy.ndarray.resize.html
+        gest_seg = gest.resize(depth_size)
+        gest_seg_np = np.array(gest_seg).astype(np.uint8)
+
+        # Match this image to its COLMAP entry by name (as in get_sparse_depth).
+        found_colmap_img = None
+        for img_id_c, col_img in colmap_rec.images.items():
+            if img_id in col_img.name:
+                found_colmap_img = col_img
+                break
+
+        point_class_names = ["apex", "eave_end_point", "flashing_end_point"]
+        edge_class_names = ["eave", "ridge", "rake", "valley", "hip", "flashing", "step_flashing", "transition_line"]
+        vertices, connections = get_vertices_and_edges_from_segmentation(
+            gest_seg_np, 
+            point_class_names,
+            edge_class_names,
+            colmap_image=found_colmap_img,
+            colmap_points3D=colmap_rec.points3D,
+            edge_th=25.0,
+            min_3d_points_for_vertex=1,
+            vertex_cluster_eps=25.0,
+            use_colmap_for_vertices=False,
+            patch_size=25
+        )
+        
+        if (len(vertices) < 2) or (len(connections) < 1):
+            print(f'Not enough vertices or connections found in image {i}, skipping.')
+            vert_edge_per_image[i] = [], [], np.empty((0, 3))
+            continue
+
+        vertices_3d = create_3d_wireframe_single_image(
+            vertices, connections, depth, colmap_rec, img_id, ade_seg
+        )
+        vert_edge_per_image[i] = vertices, connections, vertices_3d
+
+    print("Applying multi-view consistency filtering...")
+
+    total_vertices_before_filtering = sum(len(v3d) for _, _, v3d in vert_edge_per_image.values())
+    print(f"Total vertices before filtering: {total_vertices_before_filtering}")
+
+    filtered_vert_edge_per_image = filter_vertices_by_multi_view_consistency(
+        vert_edge_per_image,
+        colmap_rec,
+        good_entry['gestalt'],
+        good_entry['image_ids'],
+        gestalt_color_mapping,
+        depth_sizes,
+        min_consistent_views=1,
+        min_shared_points_for_overlap=3,
+        projection_patch_size=30
+    )
+
+    total_vertices_before = sum(len(v3d) for _, _, v3d in vert_edge_per_image.values())
+    total_vertices_after = sum(len(v3d) for _, _, v3d in filtered_vert_edge_per_image.values())
+    print(f"Multi-view filtering: {total_vertices_before} -> {total_vertices_after} vertices")
+    print(f"Filtering removed {total_vertices_before - total_vertices_after} vertices ({100*(total_vertices_before - total_vertices_after)/max(total_vertices_before,1):.1f}%)")
+
+    print("Recovering edges between filtered vertices...")
+    edges_before_recovery = sum(len(conns) for _, conns, _ in filtered_vert_edge_per_image.values())
+    
+    edge_class_names = ["eave", "ridge", "rake", "valley", "hip", "flashing", "step_flashing", "transition_line"]
+    recovered_vert_edge_per_image = recover_edges_after_vertex_filtering(
+        filtered_vert_edge_per_image,
+        good_entry,
+        edge_class_names,
+        edge_th=25.0
+    )
+    
+    edges_after_recovery = sum(len(conns) for _, conns, _ in recovered_vert_edge_per_image.values())
+    print(f"Edge recovery: {edges_before_recovery} -> {edges_after_recovery} edges")
+    
+    all_3d_vertices, connections_3d = merge_vertices_3d(recovered_vert_edge_per_image, point_class_names, 0.7)
+    all_3d_vertices_clean, connections_3d_clean  = prune_not_connected(all_3d_vertices, connections_3d, keep_largest=False)
+
+    # Merge fragmented collinear segments into continuous edges, then drop any
+    # vertices orphaned by the merge.
+    all_3d_vertices_clean, connections_3d_clean = merge_collinear_edges(all_3d_vertices_clean, connections_3d_clean, cos_threshold=-0.98)
+    all_3d_vertices_clean, connections_3d_clean = prune_not_connected(all_3d_vertices_clean, connections_3d_clean, keep_largest=False)
+
+    if (len(all_3d_vertices_clean) < 2) or len(connections_3d_clean) < 1:
+        print (f'Not enough vertices or connections in the 3D vertices')
+        return empty_solution()
+
+    return all_3d_vertices_clean, connections_3d_clean

training/edge_patch.py

@@ -0,0 +1,237 @@
+"""6D cylindrical edge-patch builder and an edge-classifier evaluation check.
+
+Provides:
+  - colmap_points_xyz_rgb / build_edge_patch_6d: build the 6D (xyz + RGB)
+    cylindrical patch around an edge from COLMAP points (radius 0.5m,
+    +0.25m extension at each end). Imported by the dataset generators.
+
+Run as a script to evaluate an edge classifier on a few samples: it scores
+each handcrafted edge, splits edges by whether they are ground-truth-positive
+(both endpoints near connected GT vertices), and compares the score
+distributions.
+"""
+import os
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
+
+import sys
+import time
+import numpy as np
+import torch
+from datasets import load_dataset
+from scipy.spatial.distance import cdist
+
+CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, CURRENT_DIR)
+
+import hc_helpers as hc
+from fast_pointnet_class import (
+    load_pointnet_model as load_pnet_class,
+    predict_class_from_patch,
+)
+from hoho2025.example_solutions import read_colmap_rec
+
+NUM_TRIALS = 3
+CYL_RADIUS = 0.5
+CYL_EXT = 0.25            # extension at each end
+GT_VERTEX_THRESH = 0.5    # how close a vertex must be to a GT vertex to count
+                          # as "matched" (= HSS vert_thresh)
+
+
+def colmap_points_xyz_rgb(colmap_rec):
+    """Return (xyz, rgb_normalized_0_1) for all COLMAP points."""
+    xyz_list, rgb_list = [], []
+    for pid, p3D in colmap_rec.points3D.items():
+        xyz_list.append(p3D.xyz)
+        rgb_list.append(p3D.color / 255.0)
+    if not xyz_list:
+        return np.empty((0, 3)), np.empty((0, 3))
+    return np.array(xyz_list), np.array(rgb_list)
+
+
+def build_edge_patch_6d(u_xyz, v_xyz, colmap_xyz, colmap_rgb):
+    """
+    Returns patch dict {'patch_6d': (M, 6)} or None if cylinder too sparse.
+    """
+    line = v_xyz - u_xyz
+    L = float(np.linalg.norm(line))
+    if L < 1e-6:
+        return None
+    direction = line / L
+
+    ext_start = u_xyz - CYL_EXT * direction
+    ext_L = L + 2 * CYL_EXT
+
+    rel = colmap_xyz - ext_start[np.newaxis, :]
+    proj = rel @ direction
+    in_bounds = (proj >= 0) & (proj <= ext_L)
+
+    closest = ext_start[np.newaxis, :] + proj[:, np.newaxis] * direction[np.newaxis, :]
+    perp = np.linalg.norm(colmap_xyz - closest, axis=1)
+    in_cyl = in_bounds & (perp <= CYL_RADIUS)
+
+    if int(in_cyl.sum()) <= 10:
+        return None
+
+    midpoint = (u_xyz + v_xyz) / 2
+    pts_centered = colmap_xyz[in_cyl] - midpoint
+    rgb_signed = colmap_rgb[in_cyl] * 2.0 - 1.0   # to [-1, 1]
+    patch_6d = np.hstack([pts_centered, rgb_signed])
+    return {'patch_6d': patch_6d}
+
+
+def label_user_edges(user_v, user_e, gt_v, gt_e, thresh=GT_VERTEX_THRESH):
+    """For each user edge, return True if it matches a GT edge.
+
+    Match := both endpoints have a nearest GT vertex within `thresh`,
+    AND those GT vertices are connected in the GT edge set.
+    """
+    if len(gt_v) == 0 or len(user_v) == 0:
+        return [None] * len(user_e)
+
+    d = cdist(user_v, gt_v)
+    user_to_gt = {}
+    for i in range(len(user_v)):
+        j = int(np.argmin(d[i]))
+        if d[i, j] < thresh:
+            user_to_gt[i] = j
+
+    gt_set = set()
+    for a, b in gt_e:
+        gt_set.add((int(min(a, b)), int(max(a, b))))
+
+    out = []
+    for u, v in user_e:
+        gu = user_to_gt.get(int(u))
+        gv = user_to_gt.get(int(v))
+        if gu is None or gv is None or gu == gv:
+            out.append(False)
+            continue
+        key = (min(gu, gv), max(gu, gv))
+        out.append(key in gt_set)
+    return out
+
+
+def smoke_test_one(sample, model, device):
+    order_id = sample['order_id']
+    print(f"\n=== {order_id} ===")
+    t0 = time.time()
+
+    try:
+        with hc.suppress_stdout():
+            user_v, user_e = hc.hc_predict(sample, {})
+    except Exception as e:
+        print(f"  user pipeline crashed: {e}")
+        return []
+    if len(user_v) == 0 or len(user_e) == 0:
+        print("  user pipeline empty")
+        return []
+    print(f"  user: {len(user_v)} vertices, {len(user_e)} edges  ({time.time()-t0:.1f}s)")
+
+    try:
+        colmap_rec = read_colmap_rec(sample['colmap'])
+    except Exception as e:
+        print(f"  colmap parse crashed: {e}")
+        return []
+
+    cm_xyz, cm_rgb = colmap_points_xyz_rgb(colmap_rec)
+    print(f"  colmap: {len(cm_xyz)} points")
+    if len(cm_xyz) == 0:
+        return []
+
+    gt_v = np.array(sample['wf_vertices']) if sample.get('wf_vertices') else np.empty((0, 3))
+    gt_e = [(int(a), int(b)) for a, b in sample.get('wf_edges', [])]
+    labels = label_user_edges(user_v, user_e, gt_v, gt_e)
+
+    results = []
+    skipped_sparse = 0
+    for (u_idx, v_idx), gt_match in zip(user_e, labels):
+        u_xyz = np.asarray(user_v[int(u_idx)])
+        v_xyz = np.asarray(user_v[int(v_idx)])
+        patch = build_edge_patch_6d(u_xyz, v_xyz, cm_xyz, cm_rgb)
+        if patch is None:
+            skipped_sparse += 1
+            continue
+        try:
+            cls_label, score = predict_class_from_patch(model, patch, device=device)
+        except Exception as e:
+            print(f"  inference crashed: {e}")
+            continue
+        results.append({
+            'order_id': order_id,
+            'u': int(u_idx),
+            'v': int(v_idx),
+            'edge_length': float(np.linalg.norm(v_xyz - u_xyz)),
+            'patch_n_pts': int(patch['patch_6d'].shape[0]),
+            'pred_label': int(cls_label) if cls_label is not None else None,
+            'score': float(score),
+            'gt_match': bool(gt_match) if gt_match is not None else None,
+        })
+
+    if skipped_sparse:
+        print(f"  {skipped_sparse} edges skipped (cylinder too sparse)")
+    n_gt_pos = sum(1 for r in results if r['gt_match'])
+    n_gt_neg = sum(1 for r in results if r['gt_match'] is False)
+    print(f"  scored {len(results)} edges: {n_gt_pos} GT-positive, "
+          f"{n_gt_neg} GT-negative")
+    return results
+
+
+def main():
+    print("Loading pnet_class.pth...")
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = load_pnet_class(
+        os.path.join(CURRENT_DIR, '..', 'pnet_class_2026.pth'), device=device)
+    print(f"  loaded on {device}")
+
+    print("\nStreaming validation...")
+    ds = load_dataset(
+        "usm3d/hoho22k_2026_trainval", split="validation",
+        streaming=True, trust_remote_code=True,
+    )
+
+    all_rows = []
+    for i, sample in enumerate(ds):
+        if i >= NUM_TRIALS:
+            break
+        all_rows.extend(smoke_test_one(sample, model, device))
+
+    if not all_rows:
+        print("\nNo results -- every sample failed.")
+        return
+
+    scores = np.array([r['score'] for r in all_rows])
+    pos_scores = np.array([r['score'] for r in all_rows if r['gt_match']])
+    neg_scores = np.array([r['score'] for r in all_rows if r['gt_match'] is False])
+
+    print(f"\n=== Aggregate over {len(all_rows)} edges "
+          f"({len({r['order_id'] for r in all_rows})} samples) ===")
+
+    print(f"\nScore distribution (all {len(scores)} edges):")
+    print(f"  mean={scores.mean():.3f}  median={np.median(scores):.3f}  "
+          f"std={scores.std():.3f}  min={scores.min():.3f}  max={scores.max():.3f}")
+    print(f"  fraction <0.05: {(scores < 0.05).mean()*100:.0f}%   "
+          f">0.65 (paper): {(scores > 0.65).mean()*100:.0f}%   "
+          f">0.99: {(scores > 0.99).mean()*100:.0f}%")
+
+    if len(pos_scores) and len(neg_scores):
+        diff = float(pos_scores.mean() - neg_scores.mean())
+        print(f"\nGT-positive ({len(pos_scores)}): "
+              f"mean={pos_scores.mean():.3f}  median={np.median(pos_scores):.3f}  "
+              f"std={pos_scores.std():.3f}")
+        print(f"GT-negative ({len(neg_scores)}): "
+              f"mean={neg_scores.mean():.3f}  median={np.median(neg_scores):.3f}  "
+              f"std={neg_scores.std():.3f}")
+        print(f"Mean delta (pos-neg): {diff:+.3f}")
+
+        # Quick AUC via ranking
+        all_pairs = sorted([(s, 1) for s in pos_scores] + [(s, 0) for s in neg_scores])
+        n_pos, n_neg = len(pos_scores), len(neg_scores)
+        rank_sum_pos = sum(rank+1 for rank, (_, lab) in enumerate(all_pairs) if lab == 1)
+        auc = (rank_sum_pos - n_pos*(n_pos+1)/2) / (n_pos * n_neg) if n_pos and n_neg else 0
+        print(f"AUC (pos > neg): {auc:.3f}")
+    else:
+        print("\nMissing pos or neg group -- can't compute discrimination.")
+
+
+if __name__ == "__main__":
+    main()

training/fast_pointnet_class.py

@@ -0,0 +1,317 @@
+"""PointNet binary classifier over 6D (xyz+rgb) point-cloud patches: model,
+dataset, training loop, and a single-patch predictor."""
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import pickle
+from torch.utils.data import Dataset, DataLoader
+from typing import List, Dict, Tuple, Optional
+import json
+
+class ClassificationPointNet(nn.Module):
+    """
+    PointNet implementation for binary classification from 6D point cloud patches.
+    Takes 6D point clouds (x,y,z,r,g,b) and predicts binary classification (edge/not edge).
+    """
+    def __init__(self, input_dim=6, max_points=1024):
+        super(ClassificationPointNet, self).__init__()
+        self.max_points = max_points
+
+        # Point-wise feature extraction.
+        self.conv1 = nn.Conv1d(input_dim, 64, 1)
+        self.conv2 = nn.Conv1d(64, 128, 1)
+        self.conv3 = nn.Conv1d(128, 256, 1)
+        self.conv4 = nn.Conv1d(256, 512, 1)
+        self.conv5 = nn.Conv1d(512, 1024, 1)
+        self.conv6 = nn.Conv1d(1024, 2048, 1)
+
+        # Classification head.
+        self.fc1 = nn.Linear(2048, 1024)
+        self.fc2 = nn.Linear(1024, 512)
+        self.fc3 = nn.Linear(512, 256)
+        self.fc4 = nn.Linear(256, 128)
+        self.fc5 = nn.Linear(128, 64)
+        self.fc6 = nn.Linear(64, 1)
+
+        self.bn1 = nn.BatchNorm1d(64)
+        self.bn2 = nn.BatchNorm1d(128)
+        self.bn3 = nn.BatchNorm1d(256)
+        self.bn4 = nn.BatchNorm1d(512)
+        self.bn5 = nn.BatchNorm1d(1024)
+        self.bn6 = nn.BatchNorm1d(2048)
+
+        self.dropout1 = nn.Dropout(0.3)
+        self.dropout2 = nn.Dropout(0.4)
+        self.dropout3 = nn.Dropout(0.5)
+        self.dropout4 = nn.Dropout(0.4)
+        self.dropout5 = nn.Dropout(0.3)
+
+    def forward(self, x):
+        """x: (B, input_dim, max_points) -> (B, 1) logits."""
+        batch_size = x.size(0)
+
+        x1 = F.relu(self.bn1(self.conv1(x)))
+        x2 = F.relu(self.bn2(self.conv2(x1)))
+        x3 = F.relu(self.bn3(self.conv3(x2)))
+        x4 = F.relu(self.bn4(self.conv4(x3)))
+        x5 = F.relu(self.bn5(self.conv5(x4)))
+        x6 = F.relu(self.bn6(self.conv6(x5)))
+
+        global_features = torch.max(x6, 2)[0]  # (B, 2048)
+
+        x = F.relu(self.fc1(global_features))
+        x = self.dropout1(x)
+        x = F.relu(self.fc2(x))
+        x = self.dropout2(x)
+        x = F.relu(self.fc3(x))
+        x = self.dropout3(x)
+        x = F.relu(self.fc4(x))
+        x = self.dropout4(x)
+        x = F.relu(self.fc5(x))
+        x = self.dropout5(x)
+        classification = self.fc6(x)
+
+        return classification
+
+class PatchClassificationDataset(Dataset):
+    """Loads saved .pkl patches for PointNet classification training."""
+
+    def __init__(self, dataset_dir: str, max_points: int = 1024, augment: bool = True):
+        self.dataset_dir = dataset_dir
+        self.max_points = max_points
+        self.augment = augment
+
+        self.patch_files = []
+        for file in os.listdir(dataset_dir):
+            if file.endswith('.pkl'):
+                self.patch_files.append(os.path.join(dataset_dir, file))
+
+        print(f"Found {len(self.patch_files)} patch files in {dataset_dir}")
+
+    def __len__(self):
+        return len(self.patch_files)
+
+    def __getitem__(self, idx):
+        """Returns (patch (6, max_points), label scalar, valid_mask (max_points,))."""
+        patch_file = self.patch_files[idx]
+
+        with open(patch_file, 'rb') as f:
+            patch_info = pickle.load(f)
+
+        patch_6d = patch_info['patch_6d']  # (N, 6)
+        label = patch_info.get('label', 0)
+
+        num_points = patch_6d.shape[0]
+
+        if num_points >= self.max_points:
+            indices = np.random.choice(num_points, self.max_points, replace=False)
+            patch_sampled = patch_6d[indices]
+            valid_mask = np.ones(self.max_points, dtype=bool)
+        else:
+            patch_sampled = np.zeros((self.max_points, 6))
+            patch_sampled[:num_points] = patch_6d
+            valid_mask = np.zeros(self.max_points, dtype=bool)
+            valid_mask[:num_points] = True
+
+        if self.augment:
+            patch_sampled = self._augment_patch(patch_sampled, valid_mask)
+
+        # conv1d wants channels first.
+        patch_tensor = torch.from_numpy(patch_sampled.T).float()  # (6, max_points)
+        label_tensor = torch.tensor(label, dtype=torch.float32)
+        valid_mask_tensor = torch.from_numpy(valid_mask)
+
+        return patch_tensor, label_tensor, valid_mask_tensor
+
+    def _augment_patch(self, patch, valid_mask):
+        """Random z-rotation, jitter, and scale on the xyz channels."""
+        valid_points = patch[valid_mask]
+
+        if len(valid_points) == 0:
+            return patch
+
+        angle = np.random.uniform(0, 2 * np.pi)
+        cos_angle = np.cos(angle)
+        sin_angle = np.sin(angle)
+        rotation_matrix = np.array([
+            [cos_angle, -sin_angle, 0],
+            [sin_angle, cos_angle, 0],
+            [0, 0, 1]
+        ])
+        valid_points[:, :3] = valid_points[:, :3] @ rotation_matrix.T
+
+        noise = np.random.normal(0, 0.01, valid_points[:, :3].shape)
+        valid_points[:, :3] += noise
+
+        scale = np.random.uniform(0.9, 1.1)
+        valid_points[:, :3] *= scale
+
+        patch[valid_mask] = valid_points
+        return patch
+
+def save_patches_dataset(patches: List[Dict], dataset_dir: str, entry_id: str):
+    """Pickle each patch to dataset_dir as {entry_id}_patch_{i}.pkl (skips existing)."""
+    os.makedirs(dataset_dir, exist_ok=True)
+
+    for i, patch in enumerate(patches):
+        filename = f"{entry_id}_patch_{i}.pkl"
+        filepath = os.path.join(dataset_dir, filename)
+        if os.path.exists(filepath):
+            continue
+        with open(filepath, 'wb') as f:
+            pickle.dump(patch, f)
+
+    print(f"Saved {len(patches)} patches for entry {entry_id}")
+
+def collate_fn(batch):
+    """Drop samples with no valid points; return None if the whole batch is empty."""
+    valid_batch = []
+    for patch_data, label, valid_mask in batch:
+        if valid_mask.sum() > 0:
+            valid_batch.append((patch_data, label, valid_mask))
+
+    if len(valid_batch) == 0:
+        return None
+
+    patch_data = torch.stack([item[0] for item in valid_batch])
+    labels = torch.stack([item[1] for item in valid_batch])
+    valid_masks = torch.stack([item[2] for item in valid_batch])
+
+    return patch_data, labels, valid_masks
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.xavier_uniform_(m.weight)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+    elif isinstance(m, nn.Linear):
+        nn.init.xavier_uniform_(m.weight)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+    elif isinstance(m, nn.BatchNorm1d):
+        nn.init.ones_(m.weight)
+        nn.init.zeros_(m.bias)
+
+def train_pointnet(dataset_dir: str, model_save_path: str, epochs: int = 100, batch_size: int = 32,
+                  lr: float = 0.001):
+    """Train ClassificationPointNet on the pickled patches in dataset_dir."""
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Training on device: {device}")
+
+    dataset = PatchClassificationDataset(dataset_dir, max_points=1024, augment=True)
+    print(f"Dataset loaded with {len(dataset)} samples")
+
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=8,
+                           collate_fn=collate_fn, drop_last=True)
+
+    model = ClassificationPointNet(input_dim=6, max_points=1024)
+    model.apply(init_weights)
+    model.to(device)
+
+    criterion = nn.BCEWithLogitsLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.5)
+
+    model.train()
+    for epoch in range(epochs):
+        total_loss = 0.0
+        correct = 0
+        total = 0
+        num_batches = 0
+
+        for batch_idx, batch_data in enumerate(dataloader):
+            if batch_data is None:
+                continue
+
+            patch_data, labels, valid_masks = batch_data
+            patch_data = patch_data.to(device)
+            labels = labels.to(device).unsqueeze(1)
+
+            optimizer.zero_grad()
+            outputs = model(patch_data)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss.item()
+            predicted = (torch.sigmoid(outputs) > 0.5).float()
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+            num_batches += 1
+
+            if batch_idx % 50 == 0:
+                print(f"Epoch {epoch+1}/{epochs}, Batch {batch_idx}, "
+                      f"Loss: {loss.item():.6f}, "
+                      f"Accuracy: {100 * correct / total:.2f}%")
+
+        avg_loss = total_loss / num_batches if num_batches > 0 else 0
+        accuracy = 100 * correct / total if total > 0 else 0
+
+        print(f"Epoch {epoch+1}/{epochs} completed, "
+              f"Avg Loss: {avg_loss:.6f}, "
+              f"Accuracy: {accuracy:.2f}%")
+
+        scheduler.step()
+
+        checkpoint_path = model_save_path.replace('.pth', f'_epoch_{epoch+1}.pth')
+        torch.save({
+            'model_state_dict': model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'epoch': epoch + 1,
+            'loss': avg_loss,
+            'accuracy': accuracy,
+        }, checkpoint_path)
+
+    torch.save({
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'epoch': epochs,
+    }, model_save_path)
+    
+    print(f"Model saved to {model_save_path}")
+    return model
+
+def load_pointnet_model(model_path: str, device: torch.device = None) -> ClassificationPointNet:
+    """Load a trained ClassificationPointNet in eval mode."""
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    model = ClassificationPointNet(input_dim=6, max_points=1024)
+
+    checkpoint = torch.load(model_path, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+
+    model.to(device)
+    model.eval()
+
+    return model
+
+def predict_class_from_patch(model: ClassificationPointNet, patch: Dict, device: torch.device = None) -> Tuple[int, float]:
+    """Score one patch (dict with 'patch_6d'). Returns (predicted_class, probability)."""
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    patch_6d = patch['patch_6d']  # (N, 6)
+
+    max_points = 1024
+    num_points = patch_6d.shape[0]
+
+    if num_points >= max_points:
+        indices = np.random.choice(num_points, max_points, replace=False)
+        patch_sampled = patch_6d[indices]
+    else:
+        patch_sampled = np.zeros((max_points, 6))
+        patch_sampled[:num_points] = patch_6d
+
+    patch_tensor = torch.from_numpy(patch_sampled.T).float().unsqueeze(0)  # (1, 6, max_points)
+    patch_tensor = patch_tensor.to(device)
+
+    with torch.no_grad():
+        outputs = model(patch_tensor)
+        probability = torch.sigmoid(outputs).item()
+        predicted_class = int(probability > 0.5)
+
+        return predicted_class, probability
+

training/gen_edge_dataset.py

@@ -0,0 +1,394 @@
+"""Generate HSS-aligned edge classifier training data from 2026 train split.
+
+Per training sample:
+  1. Run user's solution.predict_wireframe -> (user_v, user_e)
+  2. Read COLMAP points from sample['colmap']
+  3. Compute HSS_full = hss(user_v, user_e, gt_v, gt_e)
+  4. For each edge e_i in user_e:
+       Compute HSS_minus_i = hss(user_v, user_e - {e_i}, gt_v, gt_e)
+       label = 1 if HSS_full > HSS_minus_i + EPS else 0
+       (label = 1 means "removing this edge hurt HSS, so it contributes")
+       Build 6D cylinder patch (1024 x [xyz, rgb])
+  5. Save (patch, label) per edge
+
+Output: a directory of .npz files, one per sample, each containing arrays:
+  patches:   (N_edges, 1024, 6) float32  -- point cloud cylinders
+  labels:    (N_edges,) uint8            -- binary HSS-aligned labels
+  edge_meta: (N_edges, 4) int32          -- (sample_index, edge_idx, n_points, padded?)
+
+Plus a manifest .json with sample_id -> file path mapping.
+Resumable: skips samples whose output file already exists.
+"""
+import os
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
+
+import argparse
+import json
+import sys
+import time
+from collections import Counter
+
+import numpy as np
+import torch
+from datasets import load_dataset
+
+CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, CURRENT_DIR)
+
+import hc_helpers as hc
+from hoho2025.example_solutions import read_colmap_rec
+from hoho2025.metric_helper import hss as hss_fn
+from edge_patch import build_edge_patch_6d, colmap_points_xyz_rgb
+
+EPS = 1e-6           # numerical tolerance for HSS comparisons
+MAX_PATCH_POINTS = 1024
+
+
+def hss_value(v, e, gt_v, gt_e, vert_thresh=0.5, edge_thresh=0.5):
+    """Compute HSS scalar; return 0 on failure."""
+    try:
+        s = hss_fn(np.asarray(v),
+                   [(int(u), int(w)) for u, w in e],
+                   gt_v, gt_e,
+                   vert_thresh=vert_thresh, edge_thresh=edge_thresh)
+        return float(s.hss)
+    except Exception:
+        return 0.0
+
+
+def compute_per_edge_labels(user_v, user_e, gt_v, gt_e):
+    """Ablation: per-edge HSS contribution.
+
+    label[i] = 1 if removing edge i decreases HSS, else 0.
+    """
+    if len(user_e) == 0:
+        return np.zeros((0,), dtype=np.uint8)
+    full = hss_value(user_v, user_e, gt_v, gt_e)
+    labels = np.zeros(len(user_e), dtype=np.uint8)
+    for i in range(len(user_e)):
+        # Skip edge i
+        e_minus = user_e[:i] + user_e[i + 1:]
+        without = hss_value(user_v, e_minus, gt_v, gt_e)
+        labels[i] = 1 if full > without + EPS else 0
+    return labels
+
+
+def pad_or_sample_patch(patch_6d, max_pts=MAX_PATCH_POINTS, rng=None):
+    """Pad with zeros or random-sample to exactly max_pts points."""
+    if rng is None:
+        rng = np.random
+    n = patch_6d.shape[0]
+    if n >= max_pts:
+        idx = rng.choice(n, max_pts, replace=False)
+        return patch_6d[idx], False  # not padded
+    # Pad with zeros
+    out = np.zeros((max_pts, 6), dtype=np.float32)
+    out[:n] = patch_6d
+    return out, True  # padded
+
+
+def process_sample(sample, rng):
+    """Process one sample. Returns (patches, labels, edge_meta) or None."""
+    order_id = sample['order_id']
+
+    try:
+        with hc.suppress_stdout():
+            user_v, user_e = hc.hc_predict(sample, {})
+        user_v = np.asarray(user_v, dtype=np.float32)
+        user_e = [(int(u), int(w)) for u, w in user_e]
+    except Exception as e:
+        return None, f"hc_predict crashed: {e}"
+
+    if len(user_v) == 0 or len(user_e) == 0:
+        return None, "empty handcrafted output"
+
+    try:
+        colmap_rec = read_colmap_rec(sample['colmap'])
+        cm_xyz, cm_rgb = colmap_points_xyz_rgb(colmap_rec)
+    except Exception as e:
+        return None, f"colmap parse crashed: {e}"
+    if len(cm_xyz) == 0:
+        return None, "empty colmap"
+
+    # GT
+    gt_v = sample.get('wf_vertices')
+    gt_e = sample.get('wf_edges')
+    if gt_v is None or gt_e is None:
+        return None, "no GT in sample"
+    gt_v = np.asarray(gt_v, dtype=np.float32)
+    gt_e = [(int(u), int(w)) for u, w in gt_e]
+
+    # Compute per-edge HSS-aligned labels
+    labels = compute_per_edge_labels(user_v, user_e, gt_v, gt_e)
+
+    # Build per-edge patches
+    patches = np.zeros((len(user_e), MAX_PATCH_POINTS, 6), dtype=np.float32)
+    edge_meta = np.zeros((len(user_e), 4), dtype=np.int32)
+    valid_count = 0
+    valid_indices = []
+
+    for i, (u, v) in enumerate(user_e):
+        u_xyz = user_v[int(u)]
+        v_xyz = user_v[int(v)]
+        patch = build_edge_patch_6d(u_xyz, v_xyz, cm_xyz, cm_rgb)
+        if patch is None:
+            continue
+        patch_6d = patch['patch_6d'].astype(np.float32)
+        n_pts_raw = patch_6d.shape[0]
+        sampled, padded = pad_or_sample_patch(patch_6d, rng=rng)
+        patches[valid_count] = sampled
+        edge_meta[valid_count] = [i, int(u), int(v), n_pts_raw]
+        valid_indices.append(i)
+        valid_count += 1
+
+    if valid_count == 0:
+        return None, "no valid patches built"
+
+    return {
+        'patches': patches[:valid_count],
+        'labels': labels[valid_indices],
+        'edge_meta': edge_meta[:valid_count],
+        'order_id': order_id,
+        'n_edges_total': len(user_e),
+        'n_patches_valid': valid_count,
+    }, None
+
+
+def _worker_entry(sample, seed):
+    """Top-level worker function for ProcessPoolExecutor (must be picklable)."""
+    rng = np.random.RandomState(seed)
+    return process_sample(sample, rng)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--out-dir', required=True, help='Directory for .npz files')
+    parser.add_argument('--split', default='train',
+                        help='HF split: train (default) or validation')
+    parser.add_argument('--dataset', default='usm3d/hoho22k_2026_trainval:train',
+                        help='HF dataset spec. Default uses the trainval dataset '
+                             'train split (avoids dill pickling bug on Python 3.14 '
+                             'with the sampled_v2 dataset). For the canonical '
+                             'training cache, pass --dataset usm3d/s23dr-2026-sampled_4096_v2:train')
+    parser.add_argument('--max-samples', type=int, default=None,
+                        help='Limit samples (for sanity runs)')
+    parser.add_argument('--seed', type=int, default=2718)
+    parser.add_argument('--workers', type=int, default=1,
+                        help='Parallel worker processes (1=single-process; '
+                             'set to ~ncpu for large runs).')
+    parser.add_argument('--in-flight', type=int, default=None,
+                        help='Max in-flight tasks (default workers*3). Higher uses more RAM.')
+    args = parser.parse_args()
+
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    # Resolve dataset spec
+    if ':' in args.dataset:
+        repo, hf_split = args.dataset.split(':')
+    else:
+        repo, hf_split = args.dataset, args.split
+
+    print(f"Loading dataset: {repo} split={hf_split}")
+    ds = load_dataset(repo, split=hf_split, streaming=True, trust_remote_code=True)
+
+    rng = np.random.RandomState(args.seed)
+
+    manifest_path = os.path.join(args.out_dir, 'manifest.json')
+    if os.path.exists(manifest_path):
+        with open(manifest_path) as f:
+            manifest = json.load(f)
+    else:
+        manifest = {'samples': {}, 'config': {
+            'dataset': args.dataset,
+            'max_samples': args.max_samples,
+            'seed': args.seed,
+            'max_patch_points': MAX_PATCH_POINTS,
+            'workers': args.workers,
+        }}
+
+    n_done = len(manifest['samples'])
+    print(f"Resuming with {n_done} samples already done; workers={args.workers}")
+
+    # Mutable counters (closed over by handle_result)
+    state = {
+        'total_patches': 0,
+        'total_pos': 0,
+        'total_neg': 0,
+        'total_skipped': 0,
+        'seen_count': n_done,
+        't_start': time.time(),
+    }
+    crash_reasons = Counter()
+
+    def handle_result(result, err, order_id):
+        if result is None:
+            crash_reasons[err] += 1
+            manifest['samples'][order_id] = {'error': err}
+            state['total_skipped'] += 1
+        else:
+            out_path = os.path.join(args.out_dir, f'{order_id}.npz')
+            np.savez_compressed(out_path,
+                                patches=result['patches'],
+                                labels=result['labels'],
+                                edge_meta=result['edge_meta'])
+            manifest['samples'][order_id] = {
+                'path': f'{order_id}.npz',
+                'n_patches': result['n_patches_valid'],
+                'n_edges_total': result['n_edges_total'],
+            }
+            n_pos = int(result['labels'].sum())
+            n_neg = int(len(result['labels']) - n_pos)
+            state['total_patches'] += result['n_patches_valid']
+            state['total_pos'] += n_pos
+            state['total_neg'] += n_neg
+
+        state['seen_count'] += 1
+
+        if state['seen_count'] % 50 == 0:
+            with open(manifest_path, 'w') as f:
+                json.dump(manifest, f)
+
+        if state['seen_count'] % 25 == 0:
+            elapsed = time.time() - state['t_start']
+            rate = (state['seen_count'] - n_done) / max(elapsed, 1e-6)
+            pos_frac = state['total_pos'] / max(state['total_patches'], 1)
+            print(f"[{state['seen_count']}] order={order_id} "
+                  f"rate={rate:.2f}/s patches={state['total_patches']} "
+                  f"pos={state['total_pos']} neg={state['total_neg']} "
+                  f"pos_frac={pos_frac:.3f} skipped={state['total_skipped']}")
+
+    target = args.max_samples if args.max_samples else float('inf')
+
+    MAX_RETRIES = 8
+    BACKOFF = 10
+
+    if args.workers <= 1:
+        # ---- Single-process path ----
+        attempts = 0
+        stream_done = False
+        while not stream_done and state['seen_count'] < target:
+            if attempts >= MAX_RETRIES:
+                print(f"[STREAM] giving up after {attempts} attempts")
+                break
+            attempts += 1
+            if attempts > 1:
+                print(f"[STREAM] retry {attempts}/{MAX_RETRIES} after {BACKOFF}s")
+                time.sleep(BACKOFF)
+            try:
+                for sample in ds:
+                    if state['seen_count'] >= target:
+                        break
+                    order_id = sample['order_id']
+                    if order_id in manifest['samples']:
+                        continue
+                    result, err = process_sample(sample, rng)
+                    handle_result(result, err, order_id)
+                    attempts = 0
+            except Exception as e:
+                # Catch broadly: HF/webdataset streaming can raise OSError,
+                # IOError, tarfile.ReadError, EOFError, etc. on shard blips.
+                # MAX_RETRIES bounds attempts, so a real bug still surfaces.
+                print(f"[STREAM ERROR] {type(e).__name__}: {str(e)[:200]}")
+                continue
+            else:
+                stream_done = True
+                break
+    else:
+        # ---- Parallel path ----
+        from concurrent.futures import (
+            ProcessPoolExecutor, wait, FIRST_COMPLETED, as_completed,
+        )
+        import multiprocessing as mp
+        max_inflight = args.in_flight or (args.workers * 3)
+        # Force 'spawn' start method: 'fork' (Linux default) deadlocks when
+        # torch/OMP threads have been initialized in the parent.
+        ctx = mp.get_context('spawn')
+        print(f"Parallel mode: workers={args.workers}, in_flight={max_inflight}, start=spawn")
+        executor = ProcessPoolExecutor(max_workers=args.workers, mp_context=ctx)
+        pending = {}  # future -> order_id
+
+        def drain_one():
+            done, _ = wait(list(pending.keys()), return_when=FIRST_COMPLETED)
+            for f in done:
+                oid = pending.pop(f)
+                try:
+                    result, err = f.result()
+                except Exception as e:
+                    result, err = None, f"worker crashed: {type(e).__name__}: {e}"
+                handle_result(result, err, oid)
+
+        attempts = 0
+        stream_done = False
+        try:
+            while not stream_done and (state['seen_count'] + len(pending)) < target:
+                if attempts >= MAX_RETRIES:
+                    print(f"[STREAM] giving up after {attempts} attempts")
+                    break
+                attempts += 1
+                if attempts > 1:
+                    print(f"[STREAM] retry {attempts}/{MAX_RETRIES} after {BACKOFF}s")
+                    time.sleep(BACKOFF)
+                try:
+                    for sample in ds:
+                        if (state['seen_count'] + len(pending)) >= target:
+                            break
+                        order_id = sample['order_id']
+                        if order_id in manifest['samples']:
+                            continue
+                        # Throttle: wait for a worker to free up
+                        while len(pending) >= max_inflight:
+                            drain_one()
+                        seed = int(rng.randint(0, 2**31 - 1))
+                        fut = executor.submit(_worker_entry, sample, seed)
+                        pending[fut] = order_id
+                        attempts = 0
+                except Exception as e:
+                    # Catch broadly: HF/webdataset streaming can raise OSError,
+                    # IOError, tarfile.ReadError, EOFError, etc. on shard blips.
+                    # MAX_RETRIES bounds attempts, so a real bug still surfaces.
+                    print(f"[STREAM ERROR] {type(e).__name__}: {str(e)[:200]}")
+                    continue
+                else:
+                    stream_done = True
+                    break
+
+            # Drain remaining pending tasks
+            for f in as_completed(list(pending.keys())):
+                oid = pending.pop(f)
+                try:
+                    result, err = f.result()
+                except Exception as e:
+                    result, err = None, f"worker crashed: {type(e).__name__}: {e}"
+                handle_result(result, err, oid)
+        finally:
+            executor.shutdown(wait=True)
+
+    # Final save
+    with open(manifest_path, 'w') as f:
+        json.dump(manifest, f)
+
+    # Summary
+    print("\n========== Dataset gen summary ==========")
+    print(f"Samples processed: {len(manifest['samples'])}")
+    print(f"  Successful: {sum(1 for v in manifest['samples'].values() if 'path' in v)}")
+    print(f"  Skipped: {state['total_skipped']}")
+    print(f"  Crash reasons: {dict(crash_reasons.most_common(5))}")
+    print(f"\nTotal patches: {state['total_patches']}")
+    print(f"  Positive labels: {state['total_pos']} "
+          f"({100*state['total_pos']/max(state['total_patches'],1):.1f}%)")
+    print(f"  Negative labels: {state['total_neg']} "
+          f"({100*state['total_neg']/max(state['total_patches'],1):.1f}%)")
+    if state['total_patches']:
+        avg = state['total_patches'] / max(state['seen_count'], 1)
+        print(f"\nAvg patches/sample: {avg:.1f}")
+        per_patch_size = MAX_PATCH_POINTS * 6 * 4  # bytes
+        print(f"Disk estimate: "
+              f"{state['total_patches'] * per_patch_size / 1e9:.2f} GB uncompressed")
+
+    print(f"\nElapsed: {time.time()-state['t_start']:.0f}s")
+    print(f"Output: {args.out_dir}")
+    print(f"Manifest: {manifest_path}")
+
+
+if __name__ == '__main__':
+    main()

training/gen_routing_dataset.py

@@ -0,0 +1,338 @@
+"""Generate routing dataset for GBM training.
+
+Runs both pipelines per sample and collects rich per-sample features
+for a GBM that predicts which model wins (4096+seam vs 8192+noseam).
+
+Features collected:
+  Structural: n_hc, n_hc_edges, n_colmap, n_images
+  Seam diagnostics: seam_fired, seam_med_dist, seam_max_dist, n_snapped
+  Geometry: hc_spread_xy, hc_spread_z, colmap_spread_xy, colmap_spread_z
+  ML output: n_segs_4096, n_segs_8192 (surviving after conf>0.5)
+  Label: delta = sc_8192noseam - sc_4096seam
+
+Saves to a JSON file incrementally (safe to interrupt and resume). Requires the
+competition dataset (set the S23DR_DATASET environment variable to its path).
+
+Usage:
+  python3 gen_routing_dataset.py --ckpt-new path/to/8k_checkpoint.pt --n 1000
+"""
+import argparse
+import contextlib
+import io
+import json
+import os
+import sys
+import numpy as np
+import torch
+from tqdm import tqdm
+
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
+
+# Resolve modules relative to this repository.
+TRAIN_DIR = os.path.dirname(os.path.abspath(__file__))
+REPO_ROOT = os.path.abspath(os.path.join(TRAIN_DIR, '..'))
+DATASET_DIR = os.environ.get('S23DR_DATASET', 'dataset')  # competition dataset tars
+
+sys.path.insert(0, TRAIN_DIR)   # local_dataset
+sys.path.insert(0, REPO_ROOT)   # script, solution, edge_classifier, vertex_refiner
+
+import script as s
+import solution as sol
+import edge_classifier as ec
+import vertex_refiner as vr
+from local_dataset import iter_split
+from hoho2025.metric_helper import hss
+from s23dr_2026_example.tokenizer import EdgeDepthSequenceConfig
+from s23dr_2026_example.model import EdgeDepthSegmentsModel
+
+EDGE_AUG_THR = 0.55
+VERT_AUG_THR = 0.55
+SEAM_RADIUS  = 1.1
+SEAM_GUARD   = 2.0
+CONF_THRESH  = 0.5
+
+
+@contextlib.contextmanager
+def suppress_stdout():
+    saved = sys.stdout; sys.stdout = io.StringIO()
+    try: yield
+    finally: sys.stdout = saved
+
+
+def load_8192_model(ckpt_path, device):
+    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+    args = ckpt.get('args', {})
+    norm_class = torch.nn.RMSNorm if args.get('rms_norm', True) else None
+    model = EdgeDepthSegmentsModel(
+        seq_cfg=EdgeDepthSequenceConfig(seq_len=8192, colmap_points=6144, depth_points=2048),
+        segments=args.get('segments', 64), hidden=args.get('hidden', 256),
+        num_heads=args.get('num_heads', 4), kv_heads_cross=args.get('kv_heads_cross', 2),
+        kv_heads_self=args.get('kv_heads_self', 2), dim_feedforward=args.get('ff', 1024),
+        dropout=0.0, latent_tokens=args.get('latent_tokens', 256),
+        latent_layers=args.get('latent_layers', 7), decoder_layers=args.get('decoder_layers', 3),
+        cross_attn_interval=args.get('cross_attn_interval', 4), norm_class=norm_class,
+        segment_conf=args.get('segment_conf', True),
+        segment_param=args.get('segment_param', 'midpoint_dir_len'),
+        behind_emb_dim=args.get('behind_emb_dim', 8),
+        use_vote_features=args.get('vote_features', True),
+        qk_norm=args.get('qk_norm', True), qk_norm_type=args.get('qk_norm_type', 'l2'),
+    ).to(device)
+    state = {k.replace('segmenter._orig_mod.', 'segmenter.'): v for k, v in ckpt['model'].items()}
+    model.load_state_dict(state)
+    model.eval()
+    return model
+
+
+def get_conf_segs(fused, model, device):
+    """Forward pass only -- return (n_conf_segs, mean_conf, min_conf)."""
+    tokens, masks = s.build_tokens_single(fused, model, device)
+    with torch.no_grad(), torch.autocast('cuda', torch.float16, enabled=(device.type == 'cuda')):
+        out = model.forward_tokens(tokens, masks)
+    if 'conf' in out:
+        conf = torch.sigmoid(out['conf'][0].float()).cpu().numpy()
+        above = conf > CONF_THRESH
+        n = int(above.sum())
+        mean_c = float(conf.mean())
+        mean_c_above = float(conf[above].mean()) if n > 0 else 0.0
+        return n, mean_c, mean_c_above
+    return 64, 1.0, 1.0
+
+
+def snap_to_hc_with_diagnostics(ml_v, ml_e, hc_v):
+    """Like snap_to_hc but also returns diagnostics."""
+    from scipy.spatial.distance import cdist
+    diag = {'fired': False, 'med_dist': float('nan'),
+            'max_dist': float('nan'), 'n_snapped': 0}
+    if hc_v is None or len(hc_v) == 0 or len(ml_v) == 0:
+        return ml_v, list(ml_e), diag
+    dists = cdist(ml_v, hc_v)
+    min_dists = np.min(dists, axis=0)
+    diag['med_dist'] = float(np.median(min_dists))
+    diag['max_dist'] = float(np.max(min_dists))
+    if diag['med_dist'] > SEAM_GUARD:
+        return ml_v, list(ml_e), diag
+    diag['fired'] = True
+    snapped = ml_v.copy()
+    used_l, used_c = set(), set()
+    rows, cols = np.where(dists <= SEAM_RADIUS)
+    if len(rows):
+        for k in np.argsort(dists[rows, cols]):
+            i, j = int(rows[k]), int(cols[k])
+            if i not in used_l and j not in used_c:
+                snapped[i] = 0.3 * snapped[i] + 0.7 * hc_v[j]
+                used_l.add(i); used_c.add(j)
+                diag['n_snapped'] += 1
+    seen, new_e = set(), []
+    for u, v in ml_e:
+        if u == v: continue
+        key = tuple(sorted((u, v)))
+        if key not in seen:
+            seen.add(key); new_e.append((u, v))
+    return snapped, new_e, diag
+
+
+def run_augments(pv, pe, hc_v, hc_e, cm_xyz, cm_rgb, edge_model, vertex_model, device, rng):
+    pv = np.asarray(pv)
+    if edge_model is not None and hc_e and len(cm_xyz) > 0 and len(hc_v) > 0:
+        try:
+            e_sc = ec.score_edges_batched(edge_model, device, hc_v, hc_e, cm_xyz, cm_rgb, rng=rng)
+            pv, pe = ec.augment_hybrid_with_filtered_hc(pv, pe, hc_v, hc_e, e_sc, thresh=EDGE_AUG_THR)
+            pv = np.asarray(pv)
+        except Exception:
+            pass
+    if vertex_model is not None and len(cm_xyz) > 0 and len(hc_v) > 0:
+        try:
+            v_sc = vr.score_vertices_batched(vertex_model, device, hc_v, cm_xyz, cm_rgb, rng=rng)
+            pv, pe = vr.augment_hybrid_with_filtered_hc_vertices(pv, pe, hc_v, v_sc, threshold=VERT_AUG_THR)
+            pv = np.asarray(pv)
+        except Exception:
+            pass
+    return pv, [(int(a), int(b)) for a, b in pe]
+
+
+def score_hss(pv, pe, gt_v, gt_e):
+    try:
+        m = hss(np.asarray(pv).tolist(), [(int(a), int(b)) for a, b in pe],
+                gt_v, gt_e, vert_thresh=0.5, edge_thresh=0.5)
+        return float(m.hss)
+    except Exception:
+        return 0.0
+
+
+def geometry_stats(pts):
+    """Return spread stats for a point cloud (N,3)."""
+    if len(pts) < 2:
+        return {'spread_xy': 0.0, 'spread_z': 0.0, 'n': len(pts)}
+    xy_std = float(np.std(pts[:, :2]))
+    z_std  = float(np.std(pts[:, 2]))
+    return {'spread_xy': xy_std, 'spread_z': z_std, 'n': len(pts)}
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument('--ckpt-new', required=True)
+    p.add_argument('--ckpt-old', default=f'{REPO_ROOT}/checkpoint.pt')
+    p.add_argument('--split', default='validation')
+    p.add_argument('--n', type=int, default=1000)
+    p.add_argument('--out', default=f'{REPO_ROOT}/routing_dataset.json')
+    args = p.parse_args()
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Device: {device}")
+
+    model_old = s.load_model(args.ckpt_old, device)
+    model_new = load_8192_model(args.ckpt_new, device)
+    edge_model   = ec.load_pnet_class(f'{REPO_ROOT}/pnet_class_2026.pth', device=device)
+    vertex_model = vr.load_vertex_model(f'{REPO_ROOT}/vertex_refiner.pth', device=device)
+    print(f"Models loaded. Collecting {args.n} samples from '{args.split}' split.")
+    print(f"Output: {args.out}")
+
+    # Resume from existing file
+    records = []
+    seen_ids = set()
+    if os.path.exists(args.out):
+        with open(args.out) as f:
+            records = json.load(f).get('records', [])
+        seen_ids = {r['order_id'] for r in records}
+        print(f"Resuming: {len(seen_ids)} already done.")
+
+    rng = np.random.RandomState(31415)
+
+    pbar = tqdm(total=args.n, initial=len(records))
+    for sample in iter_split(DATASET_DIR, args.split):
+        if len(records) >= args.n:
+            break
+        oid = sample.get('order_id', '')
+        if oid in seen_ids:
+            continue
+        gt_v_raw = sample.get('wf_vertices')
+        if not gt_v_raw:
+            continue
+        gt_v_l = np.asarray(gt_v_raw, dtype=np.float32).tolist()
+        gt_e   = [(int(a), int(b)) for a, b in sample['wf_edges']]
+
+        n_images = len(sample.get('image_ids', []))
+
+        # HC
+        try:
+            with suppress_stdout():
+                hc_v_raw, hc_e_raw = sol.predict_wireframe(sample)
+            hc_v = np.asarray(hc_v_raw, dtype=np.float32)
+            hc_e = [(int(a), int(b)) for a, b in hc_e_raw]
+        except Exception:
+            hc_v, hc_e = np.empty((0, 3)), []
+        hc_geo = geometry_stats(hc_v)
+
+        # COLMAP
+        cm_xyz, cm_rgb = np.empty((0, 3)), np.empty((0, 3))
+        try:
+            rec = sol.read_colmap_rec(sample.get('colmap') or sample.get('colmap_binary'))
+            cm_xyz, cm_rgb = ec.colmap_points_xyz_rgb(rec)
+        except Exception:
+            pass
+        cm_geo = geometry_stats(cm_xyz)
+
+        # --- 4096 + seam + augments ---
+        s.SEQ_LEN = 4096; s.COLMAP_QUOTA = 3072; s.DEPTH_QUOTA = 1024
+        f4 = s.fuse_and_sample(sample, s.FuserConfig(), rng)
+        sc_old = 0.0
+        n4, mean_conf4, mean_conf4_above = 0, 0.0, 0.0
+        seam_diag = {'fired': False, 'med_dist': float('nan'),
+                     'max_dist': float('nan'), 'n_snapped': 0}
+        if f4 is not None:
+            n4, mean_conf4, mean_conf4_above = get_conf_segs(f4, model_old, device)
+            pv, pe = s.predict_sample(f4, model_old, device)
+            pv, pe, seam_diag = snap_to_hc_with_diagnostics(np.asarray(pv), pe, hc_v)
+            pv, pe = run_augments(pv, pe, hc_v, hc_e, cm_xyz, cm_rgb,
+                                   edge_model, vertex_model, device, rng)
+            sc_old = score_hss(pv, pe, gt_v_l, gt_e)
+
+        # --- 8192 + NO seam + augments ---
+        s.SEQ_LEN = 8192; s.COLMAP_QUOTA = 6144; s.DEPTH_QUOTA = 2048
+        f8 = s.fuse_and_sample(sample, s.FuserConfig(), rng)
+        sc_new = 0.0
+        n8, mean_conf8, mean_conf8_above = 0, 0.0, 0.0
+        if f8 is not None:
+            n8, mean_conf8, mean_conf8_above = get_conf_segs(f8, model_new, device)
+            pv, pe = s.predict_sample(f8, model_new, device)
+            pv, pe = run_augments(np.asarray(pv), pe, hc_v, hc_e, cm_xyz, cm_rgb,
+                                   edge_model, vertex_model, device, rng)
+            sc_new = score_hss(pv, pe, gt_v_l, gt_e)
+
+        s.SEQ_LEN = 4096; s.COLMAP_QUOTA = 3072; s.DEPTH_QUOTA = 1024
+
+        rec = {
+            'order_id': oid,
+            'sc_old': sc_old,
+            'sc_new': sc_new,
+            'delta': sc_new - sc_old,
+            'label': 1 if sc_new > sc_old + 1e-4 else 0,  # 1=use 8192, 0=use 4096
+            # structural
+            'n_hc': hc_geo['n'],
+            'n_hc_edges': len(hc_e),
+            'hc_spread_xy': hc_geo['spread_xy'],
+            'hc_spread_z': hc_geo['spread_z'],
+            'n_colmap': cm_geo['n'],
+            'colmap_spread_xy': cm_geo['spread_xy'],
+            'colmap_spread_z': cm_geo['spread_z'],
+            'n_images': n_images,
+            # seam diagnostics
+            'seam_fired': int(seam_diag['fired']),
+            'seam_med_dist': seam_diag['med_dist'] if not np.isnan(seam_diag['med_dist']) else -1.0,
+            'seam_max_dist': seam_diag['max_dist'] if not np.isnan(seam_diag['max_dist']) else -1.0,
+            'n_snapped': seam_diag['n_snapped'],
+            'snap_rate': seam_diag['n_snapped'] / max(n4, 1),
+            # ML confidence
+            'n_segs_4096': n4,
+            'mean_conf_4096': mean_conf4,
+            'mean_conf_4096_above': mean_conf4_above,
+            'n_segs_8192': n8,
+            'mean_conf_8192': mean_conf8,
+            'mean_conf_8192_above': mean_conf8_above,
+            'conf_delta': mean_conf8 - mean_conf4,
+            'n_segs_delta': n8 - n4,
+        }
+        records.append(rec)
+        seen_ids.add(oid)
+
+        winner = '8192' if rec['delta'] > 1e-4 else ('4096' if rec['delta'] < -1e-4 else 'tie')
+        sc_arr = np.array([r['sc_old'] for r in records])
+        tqdm.write(f"[{len(records):4d}] 4096={sc_old:.3f}  8192={sc_new:.3f}  "
+                   f"delta={rec['delta']:+.3f}  hc={hc_geo['n']:3d}  "
+                   f"seam={'Y' if seam_diag['fired'] else 'N'}({seam_diag['n_snapped']})  "
+                   f"snp={seam_diag['med_dist']:.2f}  winner={winner}  "
+                   f"run={sc_arr.mean():.4f}")
+
+        # Save every 20
+        if len(records) % 20 == 0:
+            with open(args.out, 'w') as f:
+                json.dump({'records': records}, f)
+
+        pbar.update(1)
+
+    pbar.close()
+    with open(args.out, 'w') as f:
+        json.dump({'records': records}, f)
+
+    sc_old_arr = np.array([r['sc_old'] for r in records])
+    sc_new_arr = np.array([r['sc_new'] for r in records])
+    delta_arr  = np.array([r['delta']  for r in records])
+    print(f"\nDone. {len(records)} records saved to {args.out}")
+    print(f"4096+seam: mean={sc_old_arr.mean():.4f}")
+    print(f"8192+ns:   mean={sc_new_arr.mean():.4f}")
+    print(f"8192 wins: {(delta_arr > 1e-4).sum()} / {len(records)}")
+    print(f"\nFeature correlations with delta:")
+    feat_keys = ['n_hc', 'n_hc_edges', 'hc_spread_xy', 'hc_spread_z',
+                 'n_colmap', 'colmap_spread_xy', 'n_images',
+                 'seam_fired', 'seam_med_dist', 'n_snapped', 'snap_rate',
+                 'n_segs_4096', 'mean_conf_4096', 'n_segs_8192', 'mean_conf_8192', 'conf_delta']
+    for k in feat_keys:
+        vals = np.array([r[k] for r in records], dtype=float)
+        finite = np.isfinite(vals)
+        if finite.sum() > 10:
+            corr = float(np.corrcoef(vals[finite], delta_arr[finite])[0, 1])
+            print(f"  {k:30s}: {corr:+.3f}")
+
+
+if __name__ == '__main__':
+    main()

training/gen_sampled_16384.py

@@ -0,0 +1,84 @@
+"""Generate a seq_len=16384 sampled dataset from cached_full_pcd_v2.
+
+Streams usm3d/s23dr-2026-cached_full_pcd_v2:train from HuggingFace, applies
+the priority sampler at seq_len=16384, and saves a local HF dataset
+(same format as the organizers' sampled_8192_v3: order_id + npz bytes).
+
+Usage:
+  python3 gen_sampled_16384.py --out sampled_16384
+"""
+import argparse
+import io
+import os
+import sys
+import time
+
+import numpy as np
+
+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
+sys.path.insert(0, REPO_ROOT)
+from s23dr_2026_example.make_sampled_cache import _process_sample, _priority_sample  # noqa
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument('--seq-len', type=int, default=16384)
+    p.add_argument('--out', default='sampled_16384')
+    p.add_argument('--seed', type=int, default=7)
+    p.add_argument('--n', type=int, default=0, help='0=all samples')
+    args = p.parse_args()
+
+    colmap_q = args.seq_len * 3 // 4
+    depth_q  = args.seq_len - colmap_q
+    print(f"seq_len={args.seq_len}  colmap={colmap_q}  depth={depth_q}")
+    print(f"output: {args.out}")
+
+    np.random.seed(args.seed)
+
+    from datasets import load_dataset, Dataset
+    print("Loading cached_full_pcd_v2 (streaming)...")
+    src = load_dataset('usm3d/s23dr-2026-cached_full_pcd_v2', split='train', streaming=True)
+
+    records = []
+    t0 = time.perf_counter()
+    n_done = 0
+    n_err  = 0
+
+    for raw in src:
+        if args.n > 0 and n_done >= args.n:
+            break
+
+        oid = raw['order_id']
+        try:
+            d = dict(np.load(io.BytesIO(raw['data'])))
+            result = _process_sample(d, args.seq_len, colmap_q, depth_q)
+
+            buf = io.BytesIO()
+            np.savez_compressed(buf, **result)
+            records.append({'order_id': oid, 'data': buf.getvalue()})
+            n_done += 1
+        except Exception as e:
+            print(f"  SKIP {oid}: {e}")
+            n_err += 1
+            continue
+
+        if n_done % 1000 == 0:
+            elapsed = time.perf_counter() - t0
+            rate = n_done / elapsed
+            remaining = (17500 - n_done) / rate if rate > 0 else 0
+            print(f"  {n_done} done  ({rate:.1f}/s  ~{remaining/60:.1f} min remaining)")
+
+    elapsed = time.perf_counter() - t0
+    print(f"\nProcessed {n_done} samples ({n_err} errors) in {elapsed:.0f}s")
+
+    print(f"Saving HF dataset to {args.out} ...")
+    ds = Dataset.from_list(records)
+    ds.save_to_disk(args.out)
+    print(f"Done. {len(ds)} records saved to {args.out}")
+    print(f"\nFine-tune on this dataset with "
+          f"s23dr_2026_example.train --cache-dir local://{args.out} "
+          f"--seq-len {args.seq_len}")
+
+
+if __name__ == '__main__':
+    main()

training/gen_vertex_dataset.py

@@ -0,0 +1,408 @@
+"""Generate HSS-aligned vertex refiner training data from 2026 train split.
+
+Per training sample:
+  1. Run user's solution.predict_wireframe -> (user_v, user_e)  [HC pipeline output]
+  2. Read COLMAP points from sample['colmap']
+  3. For each vertex u in user_v:
+       - Build 6D sphere patch (1024 x [xyz_centered, rgb]) around u
+       - Class label = 1 if GT vertex within HSS_VERT_THRESH (50cm), else 0
+       - Regression delta = (nearest_gt_v - u) when class=1, else zero
+       - gt_distance = nearest GT vertex distance (scalar) -- useful for analysis
+  4. Save (patches, class_labels, regression_deltas, gt_distances, vertex_meta) per sample.
+
+Candidates come from the same distribution as gen_edge_dataset.py (the
+handcrafted pipeline's outputs).
+
+Output: a directory of .npz files, one per sample, each containing arrays:
+  patches:             (N_vertices, 1024, 6) float32  -- sphere point clouds
+  class_labels:        (N_vertices,) uint8           -- HSS-aligned binary
+  regression_deltas:   (N_vertices, 3) float32       -- vector to nearest GT
+  gt_distances:        (N_vertices,) float32         -- scalar distance to GT
+  vertex_meta:         (N_vertices, 3) int32         -- (orig_idx, n_pts_raw, padded?)
+
+Plus a manifest .json with sample_id -> file path mapping.
+Resumable via manifest: skips samples already recorded.
+"""
+import os
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
+
+import argparse
+import json
+import sys
+import time
+from collections import Counter
+
+import numpy as np
+import torch  # noqa: F401  -- imported to trigger CUDA-aware init in worker
+from datasets import load_dataset
+
+CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, CURRENT_DIR)
+
+import hc_helpers as hc
+from hoho2025.example_solutions import read_colmap_rec
+from edge_patch import colmap_points_xyz_rgb
+
+HSS_VERT_THRESH = 0.5         # HSS vertex TP criterion (meters)
+MAX_PATCH_POINTS = 1024
+SPHERE_RADIUS = 1.0           # meters: feature support around each vertex
+MIN_PATCH_POINTS = 10         # skip patches with <= this many COLMAP points
+
+
+def build_vertex_patch_6d(v_xyz, colmap_xyz, colmap_rgb, radius=SPHERE_RADIUS):
+    """Sphere of COLMAP points around vertex v_xyz.
+
+    Returns (M, 6) array of [xyz_relative_to_v, rgb_signed_to_-1_+1] or None
+    if too sparse to be informative.
+    """
+    rel = colmap_xyz - v_xyz[np.newaxis, :]
+    dist = np.linalg.norm(rel, axis=1)
+    in_sphere = dist <= radius
+    n_in = int(in_sphere.sum())
+    if n_in <= MIN_PATCH_POINTS:
+        return None
+    pts_centered = rel[in_sphere]
+    rgb_signed = colmap_rgb[in_sphere] * 2.0 - 1.0
+    return np.hstack([pts_centered, rgb_signed])
+
+
+def compute_vertex_labels(user_v, gt_v, vert_thresh=HSS_VERT_THRESH):
+    """Vectorized per-vertex labels.
+
+    Returns
+      class_labels:        (N,) uint8   -- 1 if nearest GT is within vert_thresh
+      regression_deltas:   (N, 3) float32 -- gt_nearest - user_v (zeros if no GT in scene)
+      gt_distances:        (N,) float32 -- scalar distance to nearest GT (inf if no GT)
+    """
+    n = len(user_v)
+    if n == 0 or len(gt_v) == 0:
+        return (np.zeros(n, dtype=np.uint8),
+                np.zeros((n, 3), dtype=np.float32),
+                np.full(n, np.inf, dtype=np.float32))
+    # (N, G) distance matrix
+    diffs = user_v[:, np.newaxis, :] - gt_v[np.newaxis, :, :]
+    dists = np.linalg.norm(diffs, axis=2)
+    j = np.argmin(dists, axis=1)
+    gt_dists = dists[np.arange(n), j].astype(np.float32)
+    reg_deltas = (gt_v[j] - user_v).astype(np.float32)
+    class_labels = (gt_dists <= vert_thresh).astype(np.uint8)
+    return class_labels, reg_deltas, gt_dists
+
+
+def pad_or_sample_patch(patch_6d, max_pts=MAX_PATCH_POINTS, rng=None):
+    """Pad with zeros or random-sample to exactly max_pts points."""
+    if rng is None:
+        rng = np.random
+    n = patch_6d.shape[0]
+    if n >= max_pts:
+        idx = rng.choice(n, max_pts, replace=False)
+        return patch_6d[idx], False
+    out = np.zeros((max_pts, 6), dtype=np.float32)
+    out[:n] = patch_6d
+    return out, True
+
+
+def process_sample(sample, rng):
+    """One sample -> dict of arrays or (None, err_str)."""
+    order_id = sample['order_id']
+
+    # User's HC pipeline (same call as gen_edge_dataset.py -- same distribution).
+    try:
+        with hc.suppress_stdout():
+            user_v, _user_e = hc.hc_predict(sample, {})
+        user_v = np.asarray(user_v, dtype=np.float32)
+    except Exception as e:
+        return None, f"hc_predict crashed: {e}"
+
+    if len(user_v) == 0:
+        return None, "empty handcrafted vertices"
+
+    # COLMAP point cloud.
+    try:
+        colmap_rec = read_colmap_rec(sample['colmap'])
+        cm_xyz, cm_rgb = colmap_points_xyz_rgb(colmap_rec)
+    except Exception as e:
+        return None, f"colmap parse crashed: {e}"
+    if len(cm_xyz) == 0:
+        return None, "empty colmap"
+
+    # GT vertices (for HSS-aligned labels).
+    gt_v_raw = sample.get('wf_vertices')
+    if gt_v_raw is None:
+        return None, "no GT in sample"
+    gt_v = np.asarray(gt_v_raw, dtype=np.float32)
+
+    # Per-vertex labels (vectorized).
+    class_labels_all, reg_deltas_all, gt_dists_all = compute_vertex_labels(user_v, gt_v)
+
+    # Per-vertex sphere patches.
+    n = len(user_v)
+    patches = np.zeros((n, MAX_PATCH_POINTS, 6), dtype=np.float32)
+    vertex_meta = np.zeros((n, 3), dtype=np.int32)  # (orig_idx, n_pts_raw, padded_flag)
+    valid_indices = []
+    valid_count = 0
+    for i in range(n):
+        v_xyz = user_v[i]
+        raw = build_vertex_patch_6d(v_xyz, cm_xyz, cm_rgb)
+        if raw is None:
+            continue
+        n_pts_raw = raw.shape[0]
+        sampled, padded = pad_or_sample_patch(raw, rng=rng)
+        patches[valid_count] = sampled
+        vertex_meta[valid_count] = [i, n_pts_raw, int(padded)]
+        valid_indices.append(i)
+        valid_count += 1
+
+    if valid_count == 0:
+        return None, "no valid patches built"
+
+    return {
+        'patches': patches[:valid_count],
+        'class_labels': class_labels_all[valid_indices],
+        'regression_deltas': reg_deltas_all[valid_indices],
+        'gt_distances': gt_dists_all[valid_indices],
+        'vertex_meta': vertex_meta[:valid_count],
+        'order_id': order_id,
+        'n_vertices_total': n,
+        'n_patches_valid': valid_count,
+    }, None
+
+
+def _worker_entry(sample, seed):
+    """Top-level worker function for ProcessPoolExecutor (must be picklable)."""
+    rng = np.random.RandomState(seed)
+    return process_sample(sample, rng)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--out-dir', required=True, help='Directory for .npz files')
+    parser.add_argument('--split', default='train',
+                        help='HF split: train (default) or validation')
+    parser.add_argument('--dataset', default='usm3d/hoho22k_2026_trainval:train',
+                        help='HF dataset spec. Default uses the trainval dataset '
+                             'train split.')
+    parser.add_argument('--max-samples', type=int, default=None,
+                        help='Limit samples (for sanity runs)')
+    parser.add_argument('--seed', type=int, default=2718)
+    parser.add_argument('--workers', type=int, default=1,
+                        help='Parallel worker processes (1=single-process).')
+    parser.add_argument('--in-flight', type=int, default=None,
+                        help='Max in-flight tasks (default workers*3).')
+    args = parser.parse_args()
+
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    if ':' in args.dataset:
+        repo, hf_split = args.dataset.split(':')
+    else:
+        repo, hf_split = args.dataset, args.split
+
+    print(f"Loading dataset: {repo} split={hf_split}")
+    ds = load_dataset(repo, split=hf_split, streaming=True, trust_remote_code=True)
+
+    rng = np.random.RandomState(args.seed)
+
+    manifest_path = os.path.join(args.out_dir, 'manifest.json')
+    if os.path.exists(manifest_path):
+        with open(manifest_path) as f:
+            manifest = json.load(f)
+    else:
+        manifest = {'samples': {}, 'config': {
+            'dataset': args.dataset,
+            'max_samples': args.max_samples,
+            'seed': args.seed,
+            'max_patch_points': MAX_PATCH_POINTS,
+            'sphere_radius': SPHERE_RADIUS,
+            'hss_vert_thresh': HSS_VERT_THRESH,
+            'workers': args.workers,
+        }}
+
+    n_done = len(manifest['samples'])
+    print(f"Resuming with {n_done} samples already done; workers={args.workers}")
+
+    state = {
+        'total_patches': 0,
+        'total_pos_class': 0,
+        'total_neg_class': 0,
+        'sum_delta_mag_pos': 0.0,    # for mean |delta| on real vertices
+        'total_skipped': 0,
+        'seen_count': n_done,
+        't_start': time.time(),
+    }
+    crash_reasons = Counter()
+
+    def handle_result(result, err, order_id):
+        if result is None:
+            crash_reasons[err] += 1
+            manifest['samples'][order_id] = {'error': err}
+            state['total_skipped'] += 1
+        else:
+            out_path = os.path.join(args.out_dir, f'{order_id}.npz')
+            np.savez_compressed(out_path,
+                                patches=result['patches'],
+                                class_labels=result['class_labels'],
+                                regression_deltas=result['regression_deltas'],
+                                gt_distances=result['gt_distances'],
+                                vertex_meta=result['vertex_meta'])
+            manifest['samples'][order_id] = {
+                'path': f'{order_id}.npz',
+                'n_patches': result['n_patches_valid'],
+                'n_vertices_total': result['n_vertices_total'],
+                'n_class_pos': int(result['class_labels'].sum()),
+            }
+            n_pos = int(result['class_labels'].sum())
+            n_neg = int(len(result['class_labels']) - n_pos)
+            state['total_patches'] += result['n_patches_valid']
+            state['total_pos_class'] += n_pos
+            state['total_neg_class'] += n_neg
+            if n_pos > 0:
+                deltas = result['regression_deltas'][result['class_labels'] == 1]
+                state['sum_delta_mag_pos'] += float(np.linalg.norm(deltas, axis=1).sum())
+
+        state['seen_count'] += 1
+        if state['seen_count'] % 50 == 0:
+            with open(manifest_path, 'w') as f:
+                json.dump(manifest, f)
+        if state['seen_count'] % 25 == 0:
+            elapsed = time.time() - state['t_start']
+            rate = (state['seen_count'] - n_done) / max(elapsed, 1e-6)
+            pos_frac = state['total_pos_class'] / max(state['total_patches'], 1)
+            mean_delta_pos = (state['sum_delta_mag_pos']
+                              / max(state['total_pos_class'], 1))
+            print(f"[{state['seen_count']}] order={order_id} "
+                  f"rate={rate:.2f}/s patches={state['total_patches']} "
+                  f"pos={state['total_pos_class']} neg={state['total_neg_class']} "
+                  f"pos_frac={pos_frac:.3f} mean|delta|_pos={mean_delta_pos:.3f}m "
+                  f"skipped={state['total_skipped']}")
+
+    target = args.max_samples if args.max_samples else float('inf')
+
+    MAX_RETRIES = 8
+    BACKOFF = 10
+
+    if args.workers <= 1:
+        # ---- Single-process path ----
+        attempts = 0
+        stream_done = False
+        while not stream_done and state['seen_count'] < target:
+            if attempts >= MAX_RETRIES:
+                print(f"[STREAM] giving up after {attempts} attempts")
+                break
+            attempts += 1
+            if attempts > 1:
+                print(f"[STREAM] retry {attempts}/{MAX_RETRIES} after {BACKOFF}s")
+                time.sleep(BACKOFF)
+            try:
+                for sample in ds:
+                    if state['seen_count'] >= target:
+                        break
+                    order_id = sample['order_id']
+                    if order_id in manifest['samples']:
+                        continue
+                    result, err = process_sample(sample, rng)
+                    handle_result(result, err, order_id)
+                    attempts = 0
+            except Exception as e:
+                print(f"[STREAM ERROR] {type(e).__name__}: {str(e)[:200]}")
+                continue
+            else:
+                stream_done = True
+                break
+    else:
+        # ---- Parallel path ----
+        from concurrent.futures import (
+            ProcessPoolExecutor, wait, FIRST_COMPLETED, as_completed,
+        )
+        import multiprocessing as mp
+        max_inflight = args.in_flight or (args.workers * 3)
+        # Force 'spawn' start method: 'fork' deadlocks with torch+OMP on Linux.
+        ctx = mp.get_context('spawn')
+        print(f"Parallel mode: workers={args.workers}, in_flight={max_inflight}, start=spawn")
+        executor = ProcessPoolExecutor(max_workers=args.workers, mp_context=ctx)
+        pending = {}  # future -> order_id
+
+        def drain_one():
+            done, _ = wait(list(pending.keys()), return_when=FIRST_COMPLETED)
+            for f in done:
+                oid = pending.pop(f)
+                try:
+                    result, err = f.result()
+                except Exception as e:
+                    result, err = None, f"worker crashed: {type(e).__name__}: {e}"
+                handle_result(result, err, oid)
+
+        attempts = 0
+        stream_done = False
+        try:
+            while not stream_done and (state['seen_count'] + len(pending)) < target:
+                if attempts >= MAX_RETRIES:
+                    print(f"[STREAM] giving up after {attempts} attempts")
+                    break
+                attempts += 1
+                if attempts > 1:
+                    print(f"[STREAM] retry {attempts}/{MAX_RETRIES} after {BACKOFF}s")
+                    time.sleep(BACKOFF)
+                try:
+                    for sample in ds:
+                        if (state['seen_count'] + len(pending)) >= target:
+                            break
+                        order_id = sample['order_id']
+                        if order_id in manifest['samples']:
+                            continue
+                        while len(pending) >= max_inflight:
+                            drain_one()
+                        seed = int(rng.randint(0, 2**31 - 1))
+                        fut = executor.submit(_worker_entry, sample, seed)
+                        pending[fut] = order_id
+                        attempts = 0
+                except Exception as e:
+                    print(f"[STREAM ERROR] {type(e).__name__}: {str(e)[:200]}")
+                    continue
+                else:
+                    stream_done = True
+                    break
+
+            for f in as_completed(list(pending.keys())):
+                oid = pending.pop(f)
+                try:
+                    result, err = f.result()
+                except Exception as e:
+                    result, err = None, f"worker crashed: {type(e).__name__}: {e}"
+                handle_result(result, err, oid)
+        finally:
+            executor.shutdown(wait=True)
+
+    with open(manifest_path, 'w') as f:
+        json.dump(manifest, f)
+
+    # Summary
+    print("\n========== Vertex dataset gen summary ==========")
+    print(f"Samples processed: {len(manifest['samples'])}")
+    print(f"  Successful: {sum(1 for v in manifest['samples'].values() if 'path' in v)}")
+    print(f"  Skipped: {state['total_skipped']}")
+    print(f"  Crash reasons: {dict(crash_reasons.most_common(5))}")
+    print(f"\nTotal vertex patches: {state['total_patches']}")
+    print(f"  Class positive (within {HSS_VERT_THRESH}m of GT): "
+          f"{state['total_pos_class']} "
+          f"({100*state['total_pos_class']/max(state['total_patches'],1):.1f}%)")
+    print(f"  Class negative: {state['total_neg_class']} "
+          f"({100*state['total_neg_class']/max(state['total_patches'],1):.1f}%)")
+    if state['total_pos_class']:
+        mean_delta_pos = (state['sum_delta_mag_pos']
+                          / max(state['total_pos_class'], 1))
+        print(f"  Mean |delta| on positives: {mean_delta_pos:.3f}m  "
+              f"(if small, refinement has little to do)")
+    if state['total_patches']:
+        avg = state['total_patches'] / max(state['seen_count'], 1)
+        print(f"\nAvg patches/sample: {avg:.1f}")
+        per_patch_size = MAX_PATCH_POINTS * 6 * 4
+        print(f"Disk estimate: "
+              f"{state['total_patches'] * per_patch_size / 1e9:.2f} GB uncompressed")
+    print(f"\nElapsed: {time.time()-state['t_start']:.0f}s")
+    print(f"Output: {args.out_dir}")
+    print(f"Manifest: {manifest_path}")
+
+
+if __name__ == '__main__':
+    main()

training/hc_helpers.py

@@ -0,0 +1,31 @@
+"""Handcrafted-prediction helpers used by the dataset generators.
+
+Runs the handcrafted wireframe pipeline (solution.predict_wireframe) with
+stdout suppressed and returns cleaned (vertices, edges).
+"""
+import contextlib
+import io
+import sys
+
+import numpy as np
+
+from solution import predict_wireframe as my_predict
+
+
+@contextlib.contextmanager
+def suppress_stdout():
+    saved = sys.stdout
+    sys.stdout = io.StringIO()
+    try:
+        yield
+    finally:
+        sys.stdout = saved
+
+
+def hc_predict(sample, override=None):
+    """Run the handcrafted pipeline on a sample, returning (vertices, edges)."""
+    with suppress_stdout():
+        v, e = my_predict(sample)
+    v = np.asarray(v)
+    e = [(int(u), int(w)) for u, w in e]
+    return v, e

training/local_dataset.py

@@ -0,0 +1,92 @@
+"""Read the locally downloaded hoho22k_2026 webdataset tar files directly.
+
+Bypasses the datasets library entirely -- no version conflicts, no network.
+Each tar file contains per-scene entries:
+  {order_id}.{field}_{image_id}.npy   -- per-image numpy arrays (K, R, t, pose_only_in_colmap)
+  {order_id}.{ade|depth|gestalt}_{image_id}.png  -- per-image PIL images
+  {order_id}.colmap.zip               -- COLMAP reconstruction bytes
+  {order_id}.wf_vertices.npy          -- (N, 3) float32
+  {order_id}.wf_edges.npy             -- (M, 2) int64
+  {order_id}.wf_classifications.npy   -- (M,) int64
+
+Returns sample dicts compatible with predict_wireframe / convert_entry_to_human_readable.
+"""
+import glob
+import io
+import tarfile
+
+import numpy as np
+from PIL import Image
+
+
+VECTOR_FIELDS = {'K', 'R', 't', 'pose_only_in_colmap'}
+IMAGE_FIELDS   = {'ade', 'depth', 'gestalt'}
+GLOBAL_FIELDS  = {'colmap', 'wf_vertices', 'wf_edges', 'wf_classifications'}
+
+
+def _assemble_scene(order_id, raw):
+    """raw: {field_key_with_ext: bytes}"""
+    # Collect image_ids from any per-image field
+    image_ids = set()
+    for key in raw:
+        for prefix in [f'{f}_' for f in VECTOR_FIELDS | IMAGE_FIELDS]:
+            if key.startswith(prefix):
+                img_id = key[len(prefix):]
+                img_id = img_id.rsplit('.', 1)[0]  # strip extension
+                image_ids.add(img_id)
+    image_ids = sorted(image_ids)
+
+    sample = {'order_id': order_id, 'image_ids': image_ids}
+    for f in VECTOR_FIELDS | IMAGE_FIELDS:
+        sample[f] = []
+
+    for img_id in image_ids:
+        for field in VECTOR_FIELDS:
+            key = f'{field}_{img_id}.npy'
+            if key in raw:
+                sample[field].append(np.load(io.BytesIO(raw[key])))
+        for field in IMAGE_FIELDS:
+            key = f'{field}_{img_id}.png'
+            if key in raw:
+                sample[field].append(Image.open(io.BytesIO(raw[key])).copy())
+
+    # Global fields
+    if 'colmap.zip' in raw:
+        sample['colmap'] = raw['colmap.zip']
+    for field in ('wf_vertices', 'wf_edges', 'wf_classifications'):
+        key = f'{field}.npy'
+        if key in raw:
+            sample[field] = np.load(io.BytesIO(raw[key])).tolist()
+
+    return sample
+
+
+def iter_tar(tar_path):
+    """Yield assembled scene dicts from a single tar file."""
+    scenes = {}
+    with tarfile.open(tar_path, 'r') as tf:
+        for member in tf.getmembers():
+            if not member.isfile():
+                continue
+            name = member.name
+            dot = name.index('.')
+            order_id  = name[:dot]
+            field_key = name[dot + 1:]
+            f = tf.extractfile(member)
+            if f is None:
+                continue
+            if order_id not in scenes:
+                scenes[order_id] = {}
+            scenes[order_id][field_key] = f.read()
+
+    for order_id, raw in scenes.items():
+        yield _assemble_scene(order_id, raw)
+
+
+def iter_split(dataset_dir, split='validation'):
+    """Yield all scenes from a dataset split, sorted by tar file name."""
+    tars = sorted(glob.glob(f'{dataset_dir}/data/{split}/*.tar'))
+    if not tars:
+        raise FileNotFoundError(f'No tar files found in {dataset_dir}/data/{split}/')
+    for tar_path in tars:
+        yield from iter_tar(tar_path)