3sixty2cubemap/PLAN.md

# 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.