SHA256

Files

T

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

systems

2026-08-22 16:40:11 +02:00

14 KiB

Raw Permalink Blame History

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

Core skeleton — Cargo project, CLI, error, logging, project layout. ✅ done, cargo build clean
Stage 4 rig (pure) — full parity incl. self-test. ✅ done, cargo test + cargo run -- rig --self-test pass
wgpu scaffolding — GpuContext, texture readback; smoke kernel. ✅ done; used by extract + cubemap
Stage 0 video — ffmpeg CLI decode → RGBA pipe → GPU textures. ✅ done (no native binding)
Stage 1 extract — GPU metrics + quality filtering + chunk/spread selection. ✅ done, superset of Python, verified end-to-end
Stage 3 cubemap — projection kernel + mask reprojection. ✅ done, verified vs py360convert.e2c
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)
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
Parity polish — --gui-log-output line protocol, multi-video CSVs, edge flags. 🚧 core done; GUI smoke-test pending
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.

14 KiB Raw Permalink Blame History Unescape Escape

Rust + WebGPU rewrite plan — pipeline_360