# 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) β†’ `.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] ` β†’ 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.