pipeline_360 — Rust/WebGPU rewrite

An in-progress rewrite of the Python pipeline_360.py equirect→COLMAP-cubemap rig pipeline in Rust, with all heavy image math on the GPU through WebGPU (wgpu) for cross-vendor compatibility (Vulkan / Metal / DX12, and WebGPU/wasm).

Status: complete — single binary, masks default, native GUI, GPU masks. One pipeline_360 binary runs the full GPU pipeline and, with --features gui, also launches a native egui+wgpu GUI when invoked with no subcommand. The YOLO masks stage is on by default. Validated end-to-end on real 8K equirect video. cargo test passes; 0 clippy warnings in pipeline code. YOLO11x-seg via ONNX Runtime, and masks now run on the GPU (AMD MIGraphX/ROCm) and on the cube-face images directly (no equirect round-trip). All GPU stages are batched for throughput. See PLAN.md for the full design and milestone list.

Current state

Piece Status
math.rs (face rotations, quaternions, intrinsics) ✅ pure std, verified
stages/rig.rs (rig_config.json + nested tree + linking) ✅ pure std, verified
selftest.rs (mirrors Python rig --self-test, 6 groups) ✅ passes (cargo test)
CLI / main.rs (derive-based, mirrors Python commands) ✅ extract + cubemap + all + rig work
video/mod.rs (ffmpeg/ffprobe CLI decode → rawvideo RGBA) ✅ verified end-to-end
gpu/mod.rs (adapter prefers discrete GPU; batched readback) ✅ picks the dGPU (was defaulting to LowPower/iGPU)
gpu/shaders/equirect_to_cubemap.wgsl ✅ projection verified vs py360convert.e2c
gpu/shaders/sharpness.wgsl (workgroup-reduced metrics) ✅ verified end-to-end
gpu/shaders/c2e.wgsl (cubemap→equirect nearest) ✅ round-trip validated (legacy mask reprojection)
stages/extract.rs (GPU sharpness + quality + selection, batched) ✅ superset of Python; verified
stages/cubemap.rs (GPU equirect→6 faces, skip/resume, batched) ✅ verified vs py360convert.e2c
stages/masks.rs (YOLO11x-seg on cube faces; MIGraphX GPU; cache) ✅ IoU 0.96 vs ultralytics
gui.rs (native egui+wgpu GUI; in the same binary via --features gui) ✅ builds & runs

Build

There is one binary, pipeline_360. The YOLO masks stage is on by default (via ONNX Runtime). Use --no-default-features for a lighter build without masks.

cargo build                      # default: full pipeline incl. YOLO masks
./target/debug/pipeline_360 all --project-dir <proj> --video <v> --skip-masks
cargo test                       # runs the rig self-test (6 groups)
cargo run -- rig --self-test

# Lighter build without ONNX Runtime (no masks stage):
cargo build --no-default-features

# Run the GPU extract stage (needs ffmpeg/ffprobe on PATH + a GPU):
cargo run -- extract --video <video> --output-dir <out> \
    [--chunk-size 30] [--scale-width 1920] [--output-format png]
    [--min-sharpness 0] [--min-exposure 0 --max-exposure 255]
    [--min-contrast 0] [--min-motion 0 --max-motion 1e12]
    [--selection chunk|spread] [--target-count 0] [--min-gap 0]
    [--bad-frame-action keep|skip|warn] [--motion-mode frame_diff|off]

# Run the GPU cubemap stage (needs a Vulkan/Metal/DX12 device):
cargo run -- cubemap --frames-dir <equirect_dir> --output-frames-dir <out> \
    [--face-size 1600] [--skip-faces "U,D"] [--resume]
    [--masks-dir <dir> --output-masks-dir <mask_out>] [--jpeg-quality 95]

# Native GUI (egui + wgpu; cross-vendor/cross-platform) — same binary, with the `gui` feature:
cargo build --features gui
./target/debug/pipeline_360            # no subcommand -> launches the GUI
./target/debug/pipeline_360 --gui      # explicit
./target/debug/pipeline_360 all --project-dir <proj> --video <v>   # subcommand -> CLI

# Full pipeline (extract → cubemap → masks → rig), as the GUI invokes it.
# Note the order: cubemap runs *before* masks so masks run on the cube-face
# images (no wasteful equirect round-trip).
cargo run -- all --project-dir <proj> --video <v1> [<v2> …] \
    --extract-chunk-size 30 --cubemap-face-size 1600 --rig-ref-face F \
    --mask-model yolo11x-seg.onnx --mask-device rocm --mask-imgsz 1280

# YOLO masks stage (runs on the cube-face images the cubemap stage produced):
# Export once:  yolo export model=yolo11x-seg.pt format=onnx imgsz=1280 dynamic=True
cargo run -- masks --input-dir <03_cubemap/images> --output-dir <03_cubemap/masks> \
    [--mask-classes person,car,...] [--mask-model yolo11x-seg.onnx] \
    [--mask-device cpu|rocm] [--mask-imgsz 1280] [--mask-conf 0.35] \
    [--mask-dilate 6] [--mask-skip-faces ""] [--mask-empty-if-none] \
    [--mask-ep-cache <dir>]   # MIGraphX compiled-model cache dir (see below)

# `--mask-projection` / `--mask-face-size` are accepted for CLI/GUI backwards
# compatibility but are now no-ops: masks always run on cube-face images.

Decode uses the ffmpeg / ffprobe CLI (piped rawvideo), so no native FFmpeg bindings or video feature are required — just have ffmpeg on PATH. Fetching crates and building wgpu requires network access. In this environment that was granted, so cargo build / cargo test pass, and the extract + cubemap stages were verified end-to-end.

What the extract stage adds over the original Python

The original only picked the sharpest frame (variance-of-Laplacian) per fixed chunk. The Rust extract stage keeps that as the default and adds genuine quality control, useful for 360° SfM:

Flag Purpose
--min-sharpness Drop blurry frames (variance-of-Laplacian floor)
--min-exposure / --max-exposure Drop over/under-exposed frames (mean luma 0..255)
--min-contrast Drop flat/featureless frames (luma std floor)
--min-motion Parallax check — require inter-frame motion (baseline) between kept frames
--max-motion Reject jittery/too-fast frames (likely motion blur)
--selection spread Pick well-distributed frames instead of fixed chunks
--target-count N Target exactly N frames (spread mode)
--min-gap Enforce a minimum frame spacing between selected frames
`--bad-frame-action skip keep
`--motion-mode off frame_diff`

Metrics are written to _sharpness_metadata[_<stem>].csv as frame,sharpness,exposure,contrast,motion and reused on subsequent runs unless --extract-regenerate-metadata is given.

Batching

--extract-batch (default 8) and --cubemap-batch (default 4) submit a group of frames' GPU work in one command encoder and read it all back with a single device.poll — amortising per-frame readback latency and letting the GPU pipeline across frames. Output is identical to the non-batched path (validated for both stages).

Masks stage (GPU, on cube faces)

YOLO11x-seg segmentation via ONNX Runtime (--features masks). The masks stage runs after the cubemap stage and operates directly on the cube-face images it produced (03_cubemap/images/{stem}_{FACE}.ext). Rationale:

  • Masking equirect frames directly is useless (poles distort objects).
  • The old design round-tripped e2c → YOLO → c2e → equirect mask, then the cubemap stage did e2c again — two projections + a stitch-back, wasted.
  • Now the cubemap stage projects once; masks run on the square, undistorted cube faces and stay as cube-face masks (03_cubemap/masks/{stem}_{FACE}.ext.png).
  • Because the inputs are square, the only resize is a uniform scale to imgsz with no letterbox padding — so there is no crop/offset to mis-scale on the way back. (The previous cubemap path used r=imgsz/face_size but letterboxed with r=1.0 + a centre pad, producing shifted masks; that bug is gone by construction. Validated at IoU 0.96 vs ultralytics on a real cube face.)

Masks use the COLMAP convention: black = excluded. Movable objects (person, car, …) are masked out (black); the static scene is kept (white).

GPU execution. ort is built with load-dynamic, so it dlopens the system libonnxruntime at runtime (override with ORT_DYLIB_PATH) instead of downloading a CPU-only build. On AMD/ROCm systems that lib is built with the MIGraphX execution provider, so --mask-device rocm runs YOLO on the Radeon GPU. Falls back to CPU automatically if the GPU EP can't be registered.

ORT_DYLIB_PATH=/usr/lib/libonnxruntime.so LD_LIBRARY_PATH=/opt/rocm/lib \
cargo run -- masks --input-dir <cube_faces> --output-dir <masks> \
    --mask-model yolo11x-seg.onnx --mask-device rocm

Compiled-model cache (auto). MIGraphX compiles the graph to GPU kernels on the first run (slow — minutes for a large model on a new arch) and caches the compiled .mxr. The pipeline auto-sets ORT_MIGraphX_CACHE_PATH to ~/.cache/pipeline_360/migraphx (override with --mask-ep-cache <dir>), so the first run compiles+saves and every later run loads in seconds. Without a cache dir MIGraphX would crash at inference on this build, so the cache is set up for you automatically.

# first run: compiles + saves (slow)
... masks --mask-device rocm ...
# later runs: loads the cached model (fast)
... masks --mask-device rocm ...
# or use a custom cache dir:
... masks --mask-device rocm --mask-ep-cache /some/dir ...

The masks stage needs an ONNX export of the model (the .pt is not loadable in Rust): yolo export model=yolo11x-seg.pt format=onnx imgsz=1280 dynamic=True.

Optional offline verification (no cargo)

The stage-4 core is pure std, so it can also be compiled and run with plain rustc as a fallback when crates.io is unreachable:

rustc --edition 2021 --crate-name pipeline360 offline_check/main.rs \
      -o /tmp/pipeline360_selftest && /tmp/pipeline360_selftest
# -> self-test: 6 check groups passed.

offline_check/main.rs pulls the real source files via #[path] and is not part of the cargo build.

Native GUI

gui.rs is a native egui app rendered through wgpu (via eframe) — the same WebGPU layer as the pipeline, so it is cross-vendor/cross-platform (Vulkan/Metal/DX12). It is part of the same pipeline_360 binary, gated behind the gui feature — run the binary with no subcommand (or --gui) and it launches the GUI; give it a subcommand (all, extract, …) and it's the CLI:

cargo build --features gui
./target/debug/pipeline_360            # no subcommand -> GUI
./target/debug/pipeline_360 all --project-dir <proj> --video <v>   # subcommand -> CLI

The GUI does not call the GPU pipeline library in-process. Instead it shells out to the pipeline_360 CLI (itself, found beside the GUI exe, else on PATH) with --gui-log-output, streaming the [progress] / log line protocol on a background thread into a live log panel — the same protocol the Python GUI uses. This keeps the heavy GPU work in the (optionally release-built) CLI and keeps the UI responsive, and means the GUI works anywhere the binary does.

The UI exposes every stage's options as forms with the CLI defaults: project dir + multi-video pickers, start stage, skip-masks, and collapsible sections for extract (incl. all the quality/selection/batch knobs), masks, cubemap, and rig, plus Run/Stop and a live log with progress.

Building --features gui pulls a second wgpu major version (egui 0.36 uses wgpu 30; the pipeline core uses wgpu 25). Both coexist fine because the GUI process only renders with eframe's wgpu while the CLI subprocess does the compute with its own. For a smaller binary build the GUI in release mode.

Layout

See PLAN.md for the detailed design (architecture, stages, CLI parity, milestones, risks, verification strategy).

S
Description
A conversion script for creating a cubemap camera rig from a equirectangular video.
Readme GPL-3.0
208 KiB
Languages
Rust 94.6%
WGSL 5.4%