Funktioning rewrite, that is able to use gpu on a more variety of

systems
2026-08-22 16:40:11 +02:00
commit a44539b6db
24 changed files with 7078 additions and 0 deletions
@@ -0,0 +1,278 @@
+# Rust + WebGPU rewrite plan — `pipeline_360`
+
+This document is the detailed design for rewriting the Python
+`pipeline_360.py` / `pipeline_360_gui.py` pipeline in Rust, with all heavy
+image math running on the GPU through **WebGPU** (the `wgpu` crate) for
+maximum cross-vendor / cross-platform compatibility.
+
+Status legend at the top of each stage:
+- ✅ **done** — implemented and (for pure logic) verified
+- 🚧 **scaffold** — module/API in place, kernels/logic not complete
+- ⬜ **planned** — design decided, not started
+
+---
+
+## 1. Why Rust + WebGPU
+
+| Concern | Rust + `wgpu` |
+| --- | --- |
+| Parallelism | Rust's ownership model + `rayon` for CPU, and first-class GPU compute |
+| Cross-vendor GPU | `wgpu` = portable WebGPU: Vulkan (NVIDIA/AMD/Intel), Metal (Apple), DX12 (Windows), WebGPU (wasm). One codebase, one shader language (WGSL) |
+| Safety | No GC pauses, no segfaults/use-after-free, explicit resource lifetimes |
+| Distribution | Single static-ish binary; trivially parallel per-frame work |
+| Fit | Equirect→cubemap and sharpness are embarrassingly-parallel *per pixel* — ideal for a compute shader |
+
+Video **decode** stays on the CPU via FFmpeg (codecs are not something you
+hand-roll on a GPU bindless layer); we push the decoded RGBA frames to the GPU
+texture for all subsequent math. Hardware decode (Vulkan/VAAPI/NVDEC) is a
+later optimization.
+
+---
+
+## 2. Repository layout
+
+```
+rust_rewrite/
+  Cargo.toml             binary + lib, features: [video(fmpeg), masks(ort)]
+  PLAN.md                this file
+  src/
+    main.rs              entry point, dispatch, --gui-log-output
+    gui_main.rs          native egui+wgpu GUI (feature `gui`), drives the CLI
+    lib.rs               crate root, re-exports
+    cli.rs               clap parser mirroring the Python CLI
+    error.rs             Error/Cancelled types
+    logging.rs           emit / emit_progress ([progress] line protocol)
+    project.rs           project layout (01_frames…04_colmap), gather/scan
+    json.rs              minimal JSON value/writer (pure std)
+    math.rs              face rotations, quaternions, intrinsics (pure std)  ✅
+    selftest.rs          mirrors the Python rig self-test (pure std)        ✅
+    gpu/
+      mod.rs             GpuContext (device/queue/instance), texture helpers   ✅
+      shaders/
+        equirect_to_cubemap.wgsl   ✅ verified vs py360convert
+        sharpness.wgsl             🚧 used by stage 1
+    stages/
+      mod.rs             stage registry + orchestration
+      extract.rs         stage 1: GPU sharpness + frame writes (scaffold)
+      masks.rs           stage 2: YOLO segmentation (ONNX/wgpu) (scaffold)
+      cubemap.rs         stage 3: GPU equirect→cubemap projection           ✅
+      rig.rs             stage 4: rig config + nested tree (pure std)        ✅
+    video/
+      mod.rs             FFmpeg decode interface (frames → GPU)
+```
+
+---
+
+## 3. Stage-by-stage plan
+
+### Stage 0 — Media input (✅ done, via ffmpeg CLI)
+- **Decision change:** instead of fragile native `ffmpeg-next` bindings (the
+  system ships FFmpeg 9 / libavcodec 63, which `ffmpeg-next 7` would not
+  match), `video/mod.rs` shells out to the **`ffmpeg` / `ffprobe` CLI** and
+  pipes **`rawvideo` RGBA**. No native bindings, no `video` feature, no system
+  dev headers — just `ffmpeg` on PATH. Robust across FFmpeg versions.
+- `probe(path)` -> `{width, height, nb_frames, fps}` via `ffprobe -of json`.
+- `FramePipe::open(path, scale_width, probe)` decodes sequentially; analysis
+  decodes downscaled (`-vf scale=W:-1`), extraction decodes native res. Two
+  sequential decode passes mirror the Python analyze/extract split.
+- Multi-video = iterate several paths; per-video stem prefix + per-video
+  metadata CSV, exactly like Python.
+
+### Stage 1 — Extract best frames (✅ done, superset of Python; verified end-to-end)
+- **GPU kernel (`sharpness.wgsl`)**: one dispatch/frame computes, per 16×16
+  workgroup, the 4-neighbour **Laplacian** (matches `cv2.Laplacian` ksize=1)
+  and luma, reduces in shared memory to five partial sums
+  (`sum_g, sum_g2, sum_L, sum_L2, n`), and writes 5 floats/workgroup to a
+  storage buffer. The CPU sums the tiny partial buffer and derives:
+  - **sharpness** = var(L) (variance-of-Laplacian, the original metric),
+  - **exposure** = mean(luma),
+  - **contrast** = std(luma).
+  Readback is ~5 floats/workgroup instead of the full Laplacian plane.
+- **CPU motion proxy (`frame_diff`)**: a 64×32 luma thumbnail per frame;
+  mean-abs-diff vs the previous frame → the **parallax/baseline** metric.
+- **Quality filtering** ("throw away bad frames"): `--extract-min-sharpness`,
+  `--extract-min-exposure`/`--max-exposure`, `--extract-min-contrast`,
+  `--extract-min-motion`/`--max-motion`.
+- **Selection**:
+  - `chunk` (default, original behaviour): sharpest surviving frame per
+    `--extract-chunk-size` window; `--extract-bad-frame-action
+    skip|keep|warn` decides chunks with no survivors.
+  - `spread`: up to `--extract-target-count` sharp, well-distributed frames
+    (greedy by sharpness, accepting only if `--extract-min-gap` away from all
+    accepted).
+  - `--extract-min-gap` enforces a minimum frame spacing in both modes.
+- Defaults reproduce the original Python behaviour exactly (sharpest per
+  chunk, no filtering). Enriched CSV: `frame,sharpness,exposure,contrast,motion`
+  (reused unless `--extract-regenerate-metadata`).
+- Multi-video prefixes/metadata identical to Python.
+
+### Stage 2 — Masks / YOLO segmentation (✅ done, `ort` backend; verified vs Python)
+- **Backend:** ONNX Runtime via the `ort` crate (`--features masks`). The
+  YOLO11x-seg model must be an **ONNX** export (the `.pt` is not loadable in
+  Rust): `yolo export model=yolo11x-seg.pt format=onnx imgsz=1280 dynamic=True`
+  (dynamic batch + spatial dims so `--mask-imgsz` and `--mask-tile-batch` work).
+- **Decode (pure Rust):** output0 `[B,4+80+32,A]` (post-sigmoid class scores,
+  box cxcywh in imgsz space, 32 mask coeffs) + output1 protos `[B,32,mh,mw]`.
+  Per-class conf filter + greedy NMS (IoU `--mask-iou`, default 0.45), then per
+  detection `mask = sigmoid(coeffs ⊗ protos)` cropped to the box and scaled
+  back to the original/face resolution.
+- **Cubemap path (default):** e2c (GPU, the stage-3 kernel) → YOLO per face →
+  per-face binary masks (dilated) → **c2e** (new GPU nearest kernel,
+  `gpu/shaders/c2e.wgsl`, the gnomonic inverse of `xyzcube`) → equirect mask →
+  threshold(127) → invert (COLMAP black = excluded) → `<frame>.png`.
+- **Equirect path:** YOLO directly on the frame (letterboxed to `--mask-imgsz`).
+- All Python `--mask-*` flags retained; `--mask-empty-if-none` controls whether
+  frames with no targets get an all-keep mask or are left unmasked.
+- **Validation:** vs the Python stage on a real image with people — equirect
+  projection reproduces Python masks at **~0.89 IoU** (ONNX vs PyTorch + nearest
+  letterbox differences); cubemap path runs end-to-end (e2c→YOLO→c2e) and covers
+  the right regions. The `burn`/wgpu backend (true “everything on WebGPU”) is
+  left as a future optimization; `ort` is the pragmatic first backend.
+- **c2e note:** the shader uses a dominant-axis face selection + gnomonic
+  (tangent) UV — the geometrically exact inverse of our e2c `xyzcube`. It is
+  close to `py360convert.c2e` but not pixel-identical at curved pole boundaries
+  (py360convert uses a facetype LUT); for binary masks this is immaterial.
+
+### Stage 3 — Equirect → cubemap (✅ done, verified vs `py360convert.e2c`)
+- **`equirect_to_cubemap.wgsl`** is a faithful GPU port of `py360convert.e2c`:
+  - `xyzcube(face_w)` endpoint-**inclusive** `linspace(-0.5, 0.5, N)` grid —
+    NOT `(j+0.5)/N*2-1` (the earlier scaffold grid was a parity bug).
+  - per-face direction from the exact `py360convert` cube geometry,
+  - `xyz2uv` (`u=atan2(x,z)`, `v=atan2(y,hypot(x,z))`),
+  - `uv2coor` with the `-0.5` offsets,
+  - a 1-px-**padded** sampler replicating `EquirecSampler._pad` (top/bottom
+    rows rolled by `w/2`, left/right columns clamped) so the ±180° seam is
+    continuous; bilinear for images, nearest for masks.
+- **Resolved parity detail:** U/D faces use the *reversed-row* coordinate
+  (`rng_i = fliplr(y)`), which caused the earlier U/D mismatch.
+- **`stages/cubemap.rs`**: builds one wgpu pipeline once, uploads each equirect
+  as RGBA, dispatches per face to reusable storage textures, reads back
+  (256-byte `bytes_per_row` alignment handled in `gpu::read_texture_rgba`),
+  and writes `{stem}_{FACE}.{suffix}` (+ binary `.png` masks) via `image`.
+  Implements `--cubemap-face-size`, `--cubemap-skip-faces`, `--cubemap-resume`,
+  and `--cubemap-jpeg-quality` for jpg sources.
+- **Validation:** with a real GPU present, output was diffed against the
+  installed `py360convert`: bilinear faces match to ≤6/255 (float32 vs float64
+  rounding); nearest masks match **byte-for-byte**; skip/resume/naming all
+  match the Python stage.
+
+### Stage 4 — Rig config + nested tree (✅ done — pure std, verified)
+- Port of the Python `math.rs` primitives + `rig.rs` logic:
+  - `scan_cubemap_images/masks`, `detect_face_size` (PNG header read),
+  - `build_colmap_rig_config` (serde/minimal JSON),
+  - `compute_intrinsics` (linspace / half),
+  - materialise nested `{FACE}/{stem}.ext` tree (symlink/hardlink/copy).
+- Verified offline by compiling `math.rs` + `json.rs` + `rig.rs` +
+  `selftest.rs` with plain `rustc` (no network) and re-running the *same six
+  check groups* as `pipeline_360.py rig --self-test`.
+
+---
+
+## 4. CLI parity
+
+The Rust binary must accept the same commands/flags as the Python one so the
+existing Python GUI (`pipeline_360_gui.py`) can drive it unmodified, and so
+drops build up a `--gui-log-output` flag (SUPPRESS help) that toggles the
+`[progress] ` line protocol on stdout.
+
+- `all --project-dir … [--video v1 v2 …] --start-stage … [--skip-masks]`
+- `extract --video … …`, `masks …`, `cubemap …`, `rig …`, `rig --self-test`
+- Per-stage flags match Python defaults (chunk size 30, scale 1920, workers 4,
+  masks conf 0.35 / model yolo11x-seg.pt, cubemap face 1600 / skip U,D,
+  rig ref F / linspace / symlink, …).
+
+Exit codes: `0` ok, `1` error, `130` cancelled.
+
+> Note: the Python GUI currently passes `--gui-log-output` *after* `all`.
+> The Rust parser can accept it in both positions (add it to the top parser
+> and the subparsers) to avoid the argparse ordering bug we hit in Python.
+
+---
+
+## 5. Progress / logging line protocol
+
+Keep byte-for-byte compatibility with the Python GUI:
+- plain lines → appended log lines,
+- `[progress] <msg>` → replace-the-last-line in-place updates (tqdm-style).
+
+---
+
+## 6. Project layout compatibility (identical directories)
+
+Outputs must be byte-identical in structure to drop into the existing COLMAP
+workflow:
+
+```
+01_frames/  02_masks/  03_cubemap/{images,masks}/  04_colmap/{images,masks,rig_config.json}
+```
+
+Filenames, mask `.png` convention (black = excluded), and the nested
+`{FACE}/{stem}` tree all unchanged.
+
+---
+
+## 7. Dependencies (Cargo.toml)
+
+- `wgpu` — WebGPU implementation (the whole point).
+- `image` — PNG/JPEG/BMP/TIFF/WEBP pixel I/O.
+- `rayon` — CPU parallel orchestration (per-frame / per-face).
+- `clap` — CLI derivation.
+- `serde`/`serde_json` — config; fallback `json.rs` keeps core pure-std.
+- `anyhow`/`thiserror` — error handling.
+- `pollster` — block on async wgpu adapter/device setup from the sync CLI.
+- **Video decode uses the `ffmpeg`/`ffprobe` CLI** (no native binding, no
+  feature) — see stage 0.
+- `ort` (feature `masks`) or `burn`+wgpu — YOLO segmentation.
+
+---
+
+## 8. Milestones
+
+1. **Core skeleton** — Cargo project, CLI, error, logging, project layout.  ✅ done, `cargo build` clean
+2. **Stage 4 rig** (pure) — full parity incl. self-test.  ✅ done, `cargo test` + `cargo run -- rig --self-test` pass
+3. **wgpu scaffolding** — `GpuContext`, texture readback; smoke kernel.  ✅ done; used by extract + cubemap
+4. **Stage 0 video** — ffmpeg CLI decode → RGBA pipe → GPU textures.  ✅ done (no native binding)
+5. **Stage 1 extract** — GPU metrics + quality filtering + chunk/spread selection.  ✅ done, superset of Python, verified end-to-end
+6. **Stage 3 cubemap** — projection kernel + mask reprojection.  ✅ done, verified vs `py360convert.e2c`
+7. **`all` orchestration** — extract → cubemap → masks → rig on one GPU device.  ✅ done (order changed: cubemap before masks so masks run on the cube faces, no equirect round-trip; masks skipped with `--skip-masks`)
+8. **Stage masks** — `ort`/ONNX backend, YOLO on the cube-face images the cubemap stage produced (no e2c/c2e round-trip); GPU via MIGraphX with auto compiled-model cache.  ✅ done (`--features masks`), verified IoU 0.96 vs ultralytics
+9. **Parity polish** — `--gui-log-output` line protocol, multi-video CSVs, edge
+   flags.  🚧 core done; GUI smoke-test pending
+10. **GUI** — native **egui+wgpu** GUI (`gui_main.rs`, `--features gui`) that
+    drives the CLI subprocess via the `--gui-log-output` protocol.  ✅ done,
+    builds & runs cross-vendor; the Python GUI remains usable too
+
+### Batching (added)
+`--extract-batch` (8) and `--cubemap-batch` (4) submit a batch of frames' GPU
+work in one encoder and read it all back with one `device.poll`, amortising
+readback latency and overlapping the GPU pipeline across frames. Implemented in
+`gpu::readback_many` / `gpu::read_textures_rgba`; output is byte-identical to
+the non-batched path (validated for both stages).
+
+---
+
+## 9. Risks / open decisions
+
+- **wgpu readback latency** — extract keeps readback tiny (5 floats/workgroup);
+  cubemap reads back one face at a time. Both reuse persistent buffers/textures
+  across frames.
+- **Dependency/build weight** — `wgpu` pulls a large tree; the rig/selftest
+  core stays pure-std and buildable offline via `offline_check/`.
+- **YOLO backend maturity** — `burn`+wgpu is the least mature; `ort` is the
+  pragmatic first backend.
+- **System deps** — only `ffmpeg`/`ffprobe` on PATH (no dev headers); ONNX
+  Runtime (`--features masks`) needs its native library.
+- **Float parity** — GPU kernels use f32; py360convert is float64 numpy.
+  Cubemap output matches within float32 rounding (≤6/255 bilinear; nearest
+  masks byte-exact). The rig intrinsics are exact and deterministic.
+
+---
+
+## 10. Verification strategy
+
+- Offline (no network / no GPU): `math.rs + json.rs + rig.rs + selftest.rs`
+  compiled with plain `rustc`; the 6 rig-check groups must pass — mirrors
+  `pipeline_360.py rig --self-test` (✅ done in this first pass).
+- Online (user machine): `cargo test`, `cargo run -- rig --self-test`, then
+  stage parity checks against a sample video + each stage's output diffed
+  with the Python run.