5 Commits
Author SHA256 Message Date
Ben_Kosytorz 591bdaf08a More 2026-08-22 23:02:53 +02:00
Ben_Kosytorz f918231146 Fixes 2026-08-22 23:02:46 +02:00
Ben_Kosytorz 24a546acba Rewrote the readme to clearyfy the wibecoded nature of this script and
made it more concise. Fixed licencing. Added a wiki.
2026-08-22 16:54:05 +02:00
Ben_Kosytorz e5b7863320 Forgot the license 2026-08-22 16:42:47 +02:00
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
34 changed files with 7881 additions and 2155 deletions
+3 -5
View File
@@ -1,5 +1,3 @@
*
!.gitignore
!pipeline_360.py
!pipeline_360_gui.py
!requirements.txt
/target
Cargo.lock
offline_check/selftest_bin
+101
View File
@@ -0,0 +1,101 @@
[package]
name = "pipeline-360"
version = "0.1.0"
edition = "2021"
description = "360° equirect -> COLMAP cubemap rig pipeline (Rust/WebGPU rewrite)"
# The full GPL-3.0 text lives in ./LICENSE.
license = "GPL-3.0-or-later"
[lib]
name = "pipeline_360"
path = "src/lib.rs"
[[bin]]
name = "pipeline_360"
path = "src/main.rs"
[dependencies]
# GPU compute (the WebGPU layer) — Vulkan/Metal/DX12 on native, WebGPU on wasm.
wgpu = "25"
# Image I/O (PNG/JPEG) for extracted frames and cube faces.
image = { version = "0.25", default-features = false, features = ["png", "jpeg", "bmp", "tiff", "webp"] }
# CPU-side parallel orchestration (per-frame / per-face work).
rayon = "1.10"
# CLI parsing.
clap = { version = "4.5", features = ["derive"] }
# Serde for any structured config beyond the minimal `json` module.
serde = { version = "1", features = ["derive"] }
serde_json = "1"
# Block on the async wgpu device/adapter setup from a sync CLI.
pollster = "0.4"
# Errors / logging.
anyhow = "1"
thiserror = "2"
# Optional YOLO masks (ONNX Runtime). Gated behind a feature so the
# extract/cubemap/rig core builds without system deps. Video decode uses the
# `ffmpeg`/`ffprobe` CLI (piped rawvideo) instead of native bindings, so there
# is no `video` feature/dependency.
#
# `ort` is used with `load-dynamic` so it loads the *system* libonnxruntime at
# runtime (via dlopen, path overridable with `ORT_DYLIB_PATH`) instead of
# downloading a CPU-only build. On this box the system lib is 1.28 built with
# the MIGraphX execution provider (AMD ROCm GPU), so `--mask-device rocm` runs
# YOLO inference on the Radeon GPU. `migraphx` gates the `ep::MIGraphX` binding;
# `api-28` matches the system lib's API version.
[features]
# `masks` is on by default so a plain `cargo build` produces a fully-featured
# binary (YOLO segmentation via ONNX Runtime). Use `--no-default-features` for a
# lighter build without ONNX Runtime.
default = ["masks"]
masks = ["dep:ort", "dep:ndarray", "ort/load-dynamic", "ort/migraphx", "ort/api-28"]
# Native GUI (egui + wgpu). When enabled, the single `pipeline_360` binary also
# launches the GUI when invoked with no subcommand (or `--gui`).
gui = [
"dep:eframe",
"dep:egui",
"dep:egui-wgpu",
"dep:rfd",
]
# ort has no stable release yet; declare the rc so `cargo fetch` resolves it.
# `default-features = false` + `load-dynamic` avoids downloading/linking a
# CPU-only ONNX Runtime and instead dlopens the system lib (which has the
# MIGraphX GPU EP on AMD/ROCm systems). `ndarray` is needed for
# `TensorRef::from_array_view` / `try_extract_tensor`.
[dependencies.ort]
version = "2.0.0-rc"
optional = true
default-features = false
features = ["ndarray", "load-dynamic", "migraphx", "api-28"]
# NCHW tensor construction for ort inputs (only needed by the masks stage).
[dependencies.ndarray]
version = "0.17"
optional = true
# Native GUI deps (egui + wgpu + native file dialogs).
[dependencies.egui]
version = "0.36"
optional = true
[dependencies.eframe]
version = "0.36"
optional = true
default-features = false
features = ["wgpu", "wayland", "x11"]
[dependencies.egui-wgpu]
version = "0.36"
optional = true
[dependencies.rfd]
version = "0.15"
optional = true
[profile.release]
lto = true
codegen-units = 1
+278
View File
@@ -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.
+77 -2
View File
@@ -1,3 +1,78 @@
# 3sixty2cubemap
# pipeline_360
A conversion script for creating a cubemap camera rig from a equirectangular video.
Convert equirectangular 360° videos into a **COLMAP-ready cubemap camera-rig
dataset** — sharpest-frame extraction, movable-object masks, an
equirect→cubemap projection, and a `rig_config.json` + nested file tree.
This is the **Rust + WebGPU rewrite** of the original Python pipeline (now in
[`../python`](../python), kept as legacy reference).
> ⚡ **Status: fully vibecoded, and the rewrite is done.** Everything described
> here is implemented and running on real 8K test footage — GPU sharpness,
> YOLO segmentation masks on the GPU, cubemap projection, and rig generation all
> work end-to-end. It should just work; if something doesn't, the
> [Wiki](wiki/Home.md) and issue tracker are the place to look.
## What it does
```
video(s) → 01_frames (sharpest, quality-filtered)
→ 03_cubemap/images (equirect → 6 cube faces, on GPU)
→ 03_cubemap/masks (YOLO11x-seg, on GPU)
→ 04_colmap/{images,masks,rig_config.json} (nested tree for COLMAP)
```
Everything heavy is on the **GPU via WebGPU** (`wgpu`) — Vulkan / Metal / DX12,
so it's vendor-neutral. YOLO masks use the **system ONNX Runtime** with the AMD
MIGraphX provider for GPU inference (falls back to CPU). A native **egui GUI**
ships in the same binary.
## Build & run
```bash
cargo build # full pipeline (masks on by default)
cargo run -- rig --self-test # quick sanity check
# single command, all four stages (cube faces → masks on the GPU):
cargo run -- all --project-dir out --video vid1.mp4 [vid2.mp4 …] \
--mask-model yolo11x-seg.onnx --mask-device rocm
# native GUI (one binary, --features gui):
cargo build --features gui && ./target/debug/pipeline_360
```
Prerequisites: a Rust toolchain, `ffmpeg`/`ffprobe` on `PATH`, a GPU with a
WebGPU backend, and an **ONNX** export of the segmentation model
(`yolo export model=yolo11x-seg.pt format=onnx imgsz=1280 dynamic=True`).
## Features
- **Extract** — GPU sharpness/exposure/contrast scores + quality filtering,
parallax (`--min-motion`) and bad-frame rejection, chunk or spread selection,
multi-video, batched.
- **Masks** — YOLO11x-seg on the cube-face images (no wasteful round-trip),
COLMAP convention (black = excluded), GPU inference with auto cached
compiled model (first run compiles, later runs are fast).
- **Cubemap** — faithful port of `py360convert.e2c`, batched, skip/resume.
- **Rig** — `rig_config.json` + nested `{FACE}/{stem}` tree, symlink/hardlink/copy.
- **GUI** — native egui+wgpu, shells out to the CLI for the heavy work.
## Project structure
```
Cargo.toml, main.rs entry point, CLI, `all` orchestration
src/gpu/ GpuContext + WGSL shaders
src/stages/ extract, masks, cubemap, rig
src/gui.rs egui+wgpu GUI (feature `gui`)
wiki/ the documentation (setup, usage, GPU, troubleshooting)
```
## License
[GPL-3.0-or-later](LICENSE) — see [`LICENSE`](LICENSE).
## More
- **Wiki** → [`wiki/Home.md`](wiki/Home.md) — setup, CLI reference, pipeline
internals, GPU/MIGraphX setup, troubleshooting.
- **Design notes** → [`PLAN.md`](PLAN.md) — the original rewrite plan.
+42
View File
@@ -0,0 +1,42 @@
//! Offline verification driver — NOT part of the cargo build.
//!
//! Compiles the pure-`std` core modules of the real crate directly with
//! `rustc` (no network, no external deps) and runs the rig self-test, mirroring
//! `python pipeline_360.py rig --self-test`.
//!
//! Build + run:
//! rustc --edition 2021 --crate-name pipeline360 offline_check/main.rs \
//! -o /tmp/pipeline360_selftest && /tmp/pipeline360_selftest
#[path = "../src/error.rs"]
mod error;
#[path = "../src/json.rs"]
mod json;
#[path = "../src/logging.rs"]
mod logging;
#[path = "../src/math.rs"]
mod math;
#[path = "../src/project.rs"]
mod project;
#[path = "../src/stages/rig.rs"]
pub mod rig;
// selftest refers to `crate::stages::rig`; expose the top-level `rig` module under that path.
mod stages {
pub use crate::rig;
}
#[path = "../src/selftest.rs"]
mod selftest;
fn main() {
match selftest::run_selftest() {
Ok(ok) => {
println!("\nself-test: {ok} check groups passed.");
}
Err(e) => {
eprintln!("self-test FAILED: {e}");
std::process::exit(1);
}
}
}
-1585
View File
File diff suppressed because it is too large Load Diff
-542
View File
@@ -1,542 +0,0 @@
#!/usr/bin/env python3
"""
pipeline_360_gui.py -- single-form GUI for pipeline_360.py.
One form, all settings, one Run button that runs the whole pipeline
(extract -> masks -> cubemap -> rig) in sequence, streaming progress into
a log pane. Mirrors extract_sharpest_frame_gui.py's approach of shelling
out to the CLI script and parsing "[progress] "-prefixed lines for
in-place progress updates -- except this wraps all four stages instead of
one, via `pipeline_360.py all --gui-log-output ...`.
Differences from the original GUI this is modeled on:
- English-only. The original's EN/JA translation table added a fair
amount of boilerplate for a tool this much larger; happy to add
localisation back if it's wanted.
- Cross-platform Stop button. The original only force-killed via
Windows' `taskkill`, which silently wouldn't have worked on Linux.
This one kills the whole process group via SIGTERM/SIGKILL on POSIX
(and still uses taskkill on Windows), so Stop actually stops the
ProcessPoolExecutor workers the extract/cubemap stages spawn, not
just the parent CLI process.
- No PyInstaller-frozen-exe branch -- always invokes
`sys.executable pipeline_360.py all ...`. Add that branch back if you
package this as a standalone binary later.
Requires: tkinter (usually stdlib; on some Linux distros it's a separate
package, e.g. `sudo pacman -S tk` on Arch/CachyOS).
"""
import os
import queue
import signal
import subprocess
import sys
import threading
import tkinter as tk
from pathlib import Path
from tkinter import filedialog, messagebox, ttk
from typing import Optional
from pipeline_360 import PROGRESS_PREFIX, COCO_MOVABLE_DEFAULTS, FACES
APP_TITLE = "360 Photogrammetry Pipeline"
class ToolTip:
def __init__(self, widget, text: str = "") -> None:
self.widget = widget
self.text = text
self.tip_window = None
widget.bind("<Enter>", self.show_tip)
widget.bind("<Leave>", self.hide_tip)
def show_tip(self, _event=None) -> None:
if self.tip_window or not self.text:
return
x = self.widget.winfo_rootx() + 20
y = self.widget.winfo_rooty() + self.widget.winfo_height() + 8
self.tip_window = tw = tk.Toplevel(self.widget)
tw.wm_overrideredirect(True)
tw.wm_geometry(f"+{x}+{y}")
tk.Label(tw, text=self.text, justify=tk.LEFT, background="#ffffe0",
relief=tk.SOLID, borderwidth=1, font=("Segoe UI", 9), wraplength=360).pack(ipadx=4)
def hide_tip(self, _event=None) -> None:
if self.tip_window:
self.tip_window.destroy()
self.tip_window = None
# ------------------------------------------------------------------------
# Field spec: one entry per pipeline_360.py 'all' flag exposed in the form.
# kind: "int" | "float" | "str" | "choice" | "bool"
# ------------------------------------------------------------------------
class Field:
def __init__(self, flag, label, kind, default, choices=None, tip=""):
self.flag = flag # CLI flag, e.g. "--extract-chunk-size"
self.label = label
self.kind = kind
self.default = default
self.choices = choices or []
self.tip = tip
SECTIONS = [
("Extract Sharpest Frames", [
Field("--extract-chunk-size", "Chunk size", "int", 30,
tip="Pick one sharpest frame out of every N video frames."),
Field("--extract-scale-width", "Scale width (px)", "int", 1920,
tip="Resize width used only for sharpness scoring. 0 = full resolution."),
Field("--extract-workers", "Workers", "int", 4,
tip="Parallel worker processes for the sharpness pass."),
Field("--extract-output-format", "Output format", "choice", "png", choices=["png", "jpg"]),
Field("--extract-jpeg-quality", "JPEG quality", "int", 95,
tip="Only used when output format is jpg."),
Field("--extract-regenerate-metadata", "Regenerate sharpness cache", "bool", False,
tip="Ignore any existing _sharpness_metadata.csv in 01_frames and recompute it."),
]),
("Generate SfM Masks (YOLO)", [
Field("--mask-classes", "Classes to mask", "str", ",".join(COCO_MOVABLE_DEFAULTS),
tip="Comma-separated COCO class names to exclude from reconstruction."),
Field("--mask-model", "Model", "str", "yolo11x-seg.pt",
tip="Ultralytics seg model name or path."),
Field("--mask-device", "Device", "str", "cuda",
tip="cuda, cuda:0, cpu, or mps. ROCm PyTorch keeps the 'cuda' device string."),
Field("--mask-conf", "Confidence threshold", "float", 0.35),
Field("--mask-imgsz", "Inference size", "int", 1280),
Field("--mask-dilate", "Dilate (px)", "int", 6,
tip="Grows mask edges to cover segmentation-boundary error / motion blur."),
Field("--mask-projection", "Projection", "choice", "cubemap", choices=["cubemap", "equirect"],
tip="cubemap fixes YOLO missing people near the poles; equirect runs on the raw frame."),
Field("--mask-face-size", "Detection face size (px)", "int", 1024),
Field("--mask-skip-faces", "Skip faces", "str", "",
tip="Comma list of F/R/B/L/U/D to skip detection on, e.g. U,D."),
Field("--mask-batch", "Frame batch", "int", 8),
Field("--mask-tile-batch", "GPU tile batch", "int", 8,
tip="Bounds VRAM use directly for cubemap projection. Lower this first on OOM."),
Field("--mask-empty-if-none", "Write empty masks for clean frames", "bool", False,
tip="Off by default to save disk space; a missing mask already means 'unmasked'."),
]),
("Project to Cubemap", [
Field("--cubemap-face-size", "Face size (px)", "int", 1600,
tip="Final training image resolution per cube face."),
Field("--cubemap-skip-faces", "Skip faces", "str", "",
tip="Comma list of F/R/B/L/U/D to exclude from the exported dataset, e.g. U,D."),
Field("--cubemap-jpeg-quality", "JPEG quality", "int", 95),
Field("--cubemap-workers", "Workers", "int", 0, tip="0 = use all CPU cores."),
Field("--cubemap-resume", "Resume (skip already-written frames)", "bool", False),
]),
("Generate Rig Config", [
Field("--rig-ref-face", "Reference face", "choice", "F", choices=list(FACES),
tip="Defines the rig's coordinate frame. F (front) is the sane default."),
Field("--rig-intrinsics-convention", "Intrinsics convention", "choice", "linspace",
choices=["linspace", "half"],
tip="linspace is exact for py360convert output; only change this if COLMAP rejects the schema."),
Field("--rig-link-mode", "Link mode", "choice", "symlink", choices=["symlink", "hardlink", "copy"],
tip="How the flat cubemap output is relinked into COLMAP's required nested tree."),
]),
]
class PipelineGui(tk.Tk):
def __init__(self) -> None:
super().__init__()
self.title(APP_TITLE)
self.geometry("920x760")
self.minsize(760, 560)
self.project_dir = tk.StringVar()
self.video_path = tk.StringVar()
self.start_stage = tk.StringVar(value="extract")
self.skip_masks = tk.BooleanVar(value=False)
self.status_text = tk.StringVar(value="Ready")
self.field_vars: dict[str, tk.Variable] = {}
self.field_widgets: dict[str, tk.Widget] = {}
self.worker_process: Optional[subprocess.Popen] = None
self.log_thread: Optional[threading.Thread] = None
self.log_queue: "queue.Queue" = queue.Queue()
self.stop_requested = False
self.is_closing = False
self.progress_line_active = False
self._build_widgets()
self.protocol("WM_DELETE_WINDOW", self._on_close_requested)
self.after(100, self._drain_log_queue)
# -- layout -----------------------------------------------------------
def _build_widgets(self) -> None:
outer = ttk.Frame(self, padding=12)
outer.pack(fill="both", expand=True)
outer.rowconfigure(1, weight=3)
outer.rowconfigure(3, weight=2)
outer.columnconfigure(0, weight=1)
top = ttk.Frame(outer)
top.grid(row=0, column=0, sticky="ew", pady=(0, 8))
top.columnconfigure(1, weight=1)
ttk.Label(top, text="Project directory").grid(row=0, column=0, sticky="w", padx=(0, 8), pady=4)
entry = ttk.Entry(top, textvariable=self.project_dir)
entry.grid(row=0, column=1, sticky="ew", pady=4)
ttk.Button(top, text="Browse...", command=self._browse_project_dir).grid(row=0, column=2, padx=(8, 0), pady=4)
ttk.Label(top, text="Video file").grid(row=1, column=0, sticky="w", padx=(0, 8), pady=4)
self.video_entry = ttk.Entry(top, textvariable=self.video_path)
self.video_entry.grid(row=1, column=1, sticky="ew", pady=4)
self.video_button = ttk.Button(top, text="Browse...", command=self._browse_video)
self.video_button.grid(row=1, column=2, padx=(8, 0), pady=4)
ttk.Label(top, text="Start from stage").grid(row=2, column=0, sticky="w", padx=(0, 8), pady=4)
stage_combo = ttk.Combobox(top, state="readonly", values=["extract", "masks", "cubemap", "rig"],
textvariable=self.start_stage, width=14)
stage_combo.grid(row=2, column=1, sticky="w", pady=4)
stage_combo.bind("<<ComboboxSelected>>", lambda _e: self._sync_video_requirement())
ToolTip(stage_combo, "Resuming from 'masks' or later assumes 01_frames (and 02_masks) already exist "
"in the project directory from a prior run.")
self.skip_masks_check = ttk.Checkbutton(top, text="Skip masking entirely", variable=self.skip_masks)
self.skip_masks_check.grid(row=2, column=2, sticky="w", padx=(8, 0), pady=4)
# Scrollable settings area -----------------------------------------
settings_outer = ttk.Frame(outer)
settings_outer.grid(row=1, column=0, sticky="nsew")
settings_outer.rowconfigure(0, weight=1)
settings_outer.columnconfigure(0, weight=1)
canvas = tk.Canvas(settings_outer, highlightthickness=0)
vscroll = ttk.Scrollbar(settings_outer, orient="vertical", command=canvas.yview)
canvas.configure(yscrollcommand=vscroll.set)
canvas.grid(row=0, column=0, sticky="nsew")
vscroll.grid(row=0, column=1, sticky="ns")
inner = ttk.Frame(canvas)
inner_id = canvas.create_window((0, 0), window=inner, anchor="nw")
def _on_inner_configure(_event):
canvas.configure(scrollregion=canvas.bbox("all"))
def _on_canvas_configure(event):
canvas.itemconfigure(inner_id, width=event.width)
inner.bind("<Configure>", _on_inner_configure)
canvas.bind("<Configure>", _on_canvas_configure)
def _on_mousewheel(event):
delta = 1 if getattr(event, "num", None) == 5 or event.delta < 0 else -1
canvas.yview_scroll(delta, "units")
canvas.bind_all("<MouseWheel>", _on_mousewheel) # Windows / macOS
canvas.bind_all("<Button-4>", _on_mousewheel) # Linux X11 scroll up
canvas.bind_all("<Button-5>", _on_mousewheel) # Linux X11 scroll down
for title, fields in SECTIONS:
self._build_section(inner, title, fields)
# Buttons -------------------------------------------------------
actions = ttk.Frame(outer)
actions.grid(row=2, column=0, sticky="ew", pady=(8, 4))
self.run_button = ttk.Button(actions, text="Run pipeline", command=self._start_run)
self.run_button.pack(side="right")
self.stop_button = ttk.Button(actions, text="Stop", command=self._stop_run, state="disabled")
self.stop_button.pack(side="right", padx=(0, 8))
ttk.Button(actions, text="Clear log", command=self._clear_log).pack(side="right", padx=(0, 8))
ttk.Label(actions, textvariable=self.status_text).pack(side="left")
# Log -------------------------------------------------------------
log_frame = ttk.LabelFrame(outer, text="Log")
log_frame.grid(row=3, column=0, sticky="nsew")
log_frame.rowconfigure(0, weight=1)
log_frame.columnconfigure(0, weight=1)
self.log_text = tk.Text(log_frame, wrap="word", state="disabled", height=14)
self.log_text.grid(row=0, column=0, sticky="nsew")
log_scroll = ttk.Scrollbar(log_frame, orient="vertical", command=self.log_text.yview)
log_scroll.grid(row=0, column=1, sticky="ns")
self.log_text.configure(yscrollcommand=log_scroll.set)
self._sync_video_requirement()
def _build_section(self, parent, title: str, fields) -> None:
frame = ttk.LabelFrame(parent, text=title, padding=8)
frame.pack(fill="x", expand=True, padx=4, pady=6)
frame.columnconfigure(1, weight=1)
frame.columnconfigure(3, weight=1)
for i, field in enumerate(fields):
row, col_base = divmod(i, 2)
col = col_base * 2
if field.kind == "bool":
var = tk.BooleanVar(value=field.default)
widget = ttk.Checkbutton(frame, text=field.label, variable=var)
widget.grid(row=row, column=col, columnspan=2, sticky="w", padx=(0, 16), pady=3)
else:
ttk.Label(frame, text=field.label).grid(row=row, column=col, sticky="w", padx=(0, 6), pady=3)
if field.kind == "choice":
var = tk.StringVar(value=str(field.default))
widget = ttk.Combobox(frame, state="readonly", values=field.choices, textvariable=var, width=14)
else:
var = tk.StringVar(value=str(field.default))
widget = ttk.Entry(frame, textvariable=var, width=18)
widget.grid(row=row, column=col + 1, sticky="w", padx=(0, 16), pady=3)
if field.tip:
ToolTip(widget, field.tip)
self.field_vars[field.flag] = var
self.field_widgets[field.flag] = widget
# -- state sync ---------------------------------------------------------
def _sync_video_requirement(self) -> None:
needs_video = self.start_stage.get() == "extract"
state = "normal" if needs_video else "disabled"
self.video_entry.configure(state=state)
self.video_button.configure(state=state)
# -- browse dialogs -------------------------------------------------
def _browse_project_dir(self) -> None:
selected = filedialog.askdirectory(title="Select project directory")
if selected:
self.project_dir.set(selected)
def _browse_video(self) -> None:
selected = filedialog.askopenfilename(title="Select video file")
if selected:
self.video_path.set(selected)
# -- log helpers (mirrors extract_sharpest_frame_gui.py's in-place
# progress-line replacement so tqdm-style updates don't spam the log) --
def _clear_log(self) -> None:
self.log_text.configure(state="normal")
self.log_text.delete("1.0", tk.END)
if self.progress_line_active:
self.log_text.mark_unset("progress_start")
self.log_text.mark_unset("progress_end")
self.log_text.configure(state="disabled")
self.progress_line_active = False
def _append_log(self, message: str, replace: bool = False) -> None:
self.log_text.configure(state="normal")
if replace:
if self.progress_line_active:
self.log_text.delete("progress_start", "progress_end")
else:
self.log_text.mark_set("progress_start", "end-1c")
self.log_text.mark_gravity("progress_start", tk.LEFT)
self.log_text.insert("progress_start", message + "\n")
self.log_text.mark_set("progress_end", "progress_start +1 lines")
self.progress_line_active = True
else:
if self.progress_line_active:
self.log_text.mark_unset("progress_start")
self.log_text.mark_unset("progress_end")
self.progress_line_active = False
self.log_text.insert(tk.END, message + "\n")
self.log_text.see(tk.END)
self.log_text.configure(state="disabled")
# -- command building / validation -----------------------------------
def _validate_and_build_command(self) -> Optional[list]:
project_dir = self.project_dir.get().strip()
if not project_dir:
messagebox.showerror(APP_TITLE, "Select a project directory.")
return None
start_stage = self.start_stage.get()
video = self.video_path.get().strip()
if start_stage == "extract" and not video:
messagebox.showerror(APP_TITLE, "Select a video file, or choose a later start stage.")
return None
base_dir = Path(__file__).parent
command = [sys.executable, str(base_dir / "pipeline_360.py"), "all", "--gui-log-output",
"--project-dir", project_dir, "--start-stage", start_stage]
if start_stage == "extract":
command += ["--video", video]
if self.skip_masks.get():
command += ["--skip-masks"]
for _title, fields in SECTIONS:
for field in fields:
var = self.field_vars[field.flag]
if field.kind == "bool":
if var.get():
command.append(field.flag)
continue
raw = var.get().strip() if isinstance(var, tk.StringVar) else str(var.get())
if field.kind == "int":
try:
int(raw)
except ValueError:
messagebox.showerror(APP_TITLE, f"'{field.label}' must be a whole number (got {raw!r}).")
return None
elif field.kind == "float":
try:
float(raw)
except ValueError:
messagebox.showerror(APP_TITLE, f"'{field.label}' must be a number (got {raw!r}).")
return None
command += [field.flag, raw]
return command
# -- process lifecycle ------------------------------------------------
def _is_worker_running(self) -> bool:
return self.worker_process is not None and self.worker_process.poll() is None
def _spawn(self, command: list) -> subprocess.Popen:
kwargs = {}
if sys.platform == "win32":
kwargs["creationflags"] = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
else:
kwargs["start_new_session"] = True # own process group, so we can kill it (and its children) together
return subprocess.Popen(
command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
text=True, encoding="utf-8", errors="replace", bufsize=1, **kwargs,
)
def _force_kill_process_tree(self, process: subprocess.Popen) -> None:
if sys.platform == "win32":
subprocess.run(["taskkill", "/PID", str(process.pid), "/T", "/F"], capture_output=True, text=True, check=False)
try:
process.wait(timeout=2.0)
except subprocess.TimeoutExpired:
process.kill()
else:
try:
pgid = os.getpgid(process.pid)
os.killpg(pgid, signal.SIGTERM)
process.wait(timeout=2.0)
except (ProcessLookupError, subprocess.TimeoutExpired):
try:
os.killpg(os.getpgid(process.pid), signal.SIGKILL)
except (ProcessLookupError, PermissionError):
pass
except PermissionError:
pass
try:
process.wait(timeout=1.0)
except subprocess.TimeoutExpired:
pass
def _stream_process_output(self, process: subprocess.Popen) -> None:
assert process.stdout is not None
for line in process.stdout:
self.log_queue.put(("log", line.rstrip("\r\n")))
process.stdout.close()
return_code = process.wait()
if self.stop_requested:
self.log_queue.put(("cancelled", None))
elif return_code == 0:
self.log_queue.put(("done", None))
else:
self.log_queue.put(("error", f"Pipeline exited with code {return_code}. See log above for details."))
def _set_running_state(self, running: bool) -> None:
state = "disabled" if running else "normal"
self.run_button.configure(state=state)
for widget in self.field_widgets.values():
widget.configure(state=("disabled" if running else ("readonly" if isinstance(widget, ttk.Combobox) else "normal")))
self.skip_masks_check.configure(state=state)
self.stop_button.configure(state="normal" if running else "disabled")
self.status_text.set("Running..." if running else "Ready")
if not running:
self._sync_video_requirement()
def _start_run(self) -> None:
if self._is_worker_running():
return
command = self._validate_and_build_command()
if command is None:
return
self.stop_requested = False
self._set_running_state(True)
self._append_log(f"Running: {' '.join(command)}")
try:
self.worker_process = self._spawn(command)
except OSError as exc:
self._set_running_state(False)
messagebox.showerror(APP_TITLE, f"Could not start pipeline_360.py: {exc}")
return
self.log_thread = threading.Thread(target=self._stream_process_output, args=(self.worker_process,), daemon=True)
self.log_thread.start()
def _stop_run(self) -> None:
if not self._is_worker_running():
return
self.status_text.set("Stopping...")
self._append_log("Stopping...")
self.stop_button.configure(state="disabled")
self.stop_requested = True
self._force_kill_process_tree(self.worker_process)
self.worker_process = None
self.log_thread = None
self._set_running_state(False)
def _on_close_requested(self) -> None:
if self._is_worker_running():
if not messagebox.askyesno(APP_TITLE, "The pipeline is still running. Force stop it and close the window?"):
return
self.status_text.set("Stopping...")
self._append_log("Stopping...")
self._force_kill_process_tree(self.worker_process)
self.is_closing = True
self.destroy()
def _drain_log_queue(self) -> None:
if self.is_closing:
return
while True:
try:
kind, payload = self.log_queue.get_nowait()
except queue.Empty:
break
if kind == "log":
if payload.startswith(PROGRESS_PREFIX):
self._append_log(payload[len(PROGRESS_PREFIX):], replace=True)
else:
self._append_log(payload)
elif kind == "done":
self._append_log("Done: pipeline completed successfully.")
self.worker_process = None
self.log_thread = None
self._set_running_state(False)
messagebox.showinfo(APP_TITLE, "Pipeline completed successfully.")
elif kind == "cancelled":
self._append_log("Cancelled.")
self.worker_process = None
self.log_thread = None
self._set_running_state(False)
elif kind == "error":
self._append_log(f"Error: {payload}")
self.worker_process = None
self.log_thread = None
self._set_running_state(False)
messagebox.showerror(APP_TITLE, payload)
self.after(100, self._drain_log_queue)
def main() -> None:
app = PipelineGui()
app.mainloop()
if __name__ == "__main__":
main()
-21
View File
@@ -1,21 +0,0 @@
# Core pipeline (extract, cubemap, rig stages) -- always needed
opencv-python>=4.8
numpy>=1.24
tqdm>=4.65
py360convert>=1.0
# Masks stage only (YOLO segmentation) -- skip these if you always run
# with --skip-masks
ultralytics>=8.3
# NOTE on ultralytics/torch + ROCm:
# `pip install ultralytics` will pull in a CUDA-only torch wheel as a
# dependency if torch isn't already installed. On your RX 9070 XT setup,
# install the ROCm build of torch FIRST (from https://pytorch.org, ROCm
# variant), THEN install ultralytics, so it detects torch is already
# satisfied and doesn't overwrite it with the CUDA build. Once ROCm torch
# is in place, --mask-device cuda keeps working as-is (ROCm torch keeps
# the "cuda" device string).
# tkinter (needed for pipeline_360_gui.py) is not pip-installable -- it's
# a system package. On CachyOS/Arch: sudo pacman -S tk
+60
View File
@@ -0,0 +1,60 @@
//! Error / cancellation types (pure `std`).
/// Pipeline error. Kept lightweight (no `thiserror`) so the core stays pure.
#[derive(Debug)]
pub struct PipelineError(pub String);
impl std::fmt::Display for PipelineError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
impl std::error::Error for PipelineError {}
impl From<String> for PipelineError {
fn from(s: String) -> Self {
PipelineError(s)
}
}
impl From<&str> for PipelineError {
fn from(s: &str) -> Self {
PipelineError(s.to_string())
}
}
/// Cancellation signal, raised when the GUI/cli requests a stop.
#[derive(Debug)]
pub struct PipelineCancelled;
/// A shared cancellation flag.
#[derive(Debug, Default)]
pub struct CancelFlag(std::sync::Arc<std::sync::atomic::AtomicBool>);
impl CancelFlag {
pub fn new() -> Self {
Self::default()
}
pub fn cancel(&self) {
self.0.store(true, std::sync::atomic::Ordering::SeqCst);
}
pub fn is_cancelled(&self) -> bool {
self.0.load(std::sync::atomic::Ordering::SeqCst)
}
/// Raise `PipelineCancelled` if a stop was requested.
pub fn check(&self) -> Result<(), PipelineCancelled> {
if self.is_cancelled() {
Err(PipelineCancelled)
} else {
Ok(())
}
}
pub fn clone_flag(&self) -> CancelFlag {
CancelFlag(self.0.clone())
}
}
+470
View File
@@ -0,0 +1,470 @@
//! WebGPU (`wgpu`) compute layer.
//!
//! `GpuContext` owns the wgpu instance/device/queue so every GPU stage
//! (extract sharpness, cubemap projection, later mask inference) shares one
//! device and one queue. `wgpu` maps to Vulkan/Metal/DX12 on native and WebGPU
//! in the browser — one WGSL codebase, cross-vendor.
//!
//! Added for stage 3 (cubemap): texture upload, storage-texture output,
//! and a `read_texture_rgba` readback that correctly handles the WebGPU
//! 256-byte `bytes_per_row` alignment requirement.
use wgpu::util::DeviceExt;
/// A ready-to-use wgpu device + queue plus the instance/adapter used to get it.
pub struct GpuContext {
pub device: wgpu::Device,
pub queue: wgpu::Queue,
}
/// WebGPU requires `bytes_per_row` to be a multiple of this for
/// buffer<->texture copies.
pub const COPY_BYTES_PER_ROW_ALIGNMENT: u64 = 256;
/// Round `n` up to the nearest multiple of `256` (WebGPU copy alignment).
pub fn align_bytes_per_row(n: u64) -> u64 {
n.div_ceil(COPY_BYTES_PER_ROW_ALIGNMENT) * COPY_BYTES_PER_ROW_ALIGNMENT
}
impl GpuContext {
/// Request a compute-capable adapter, preferring the *discrete* GPU.
///
/// `RequestAdapterOptions::default()` uses `power_preference: LowPower`, which
/// on a dual-GPU system (integrated + discrete) often selects the weak iGPU
/// or a software (llvmpipe) adapter — making every GPU stage look like it
/// "isn't using the GPU". We explicitly ask for `HighPerformance` and, as a
/// fallback, enumerate all adapters and pick the first discrete one.
pub async fn new(label: &str) -> anyhow::Result<GpuContext> {
let instance = wgpu::Instance::default();
let adapter = instance
.request_adapter(&wgpu::RequestAdapterOptions {
power_preference: wgpu::PowerPreference::HighPerformance,
force_fallback_adapter: false,
compatible_surface: None,
})
.await
.or_else(|_| {
// enumerate_adapter is blocking-ish; do it on a blocking context.
// wgpu exposes enumerate_adapters synchronously.
let mut best: Option<wgpu::Adapter> = None;
for a in instance.enumerate_adapters(wgpu::Backends::all()) {
let info = a.get_info();
if info.device_type == wgpu::DeviceType::DiscreteGpu {
best = Some(a);
break;
}
if best.is_none() {
best = Some(a);
}
}
best.ok_or_else(|| {
anyhow::anyhow!("no wgpu adapter found (no Vulkan/Metal/DX12 device?)")
})
})?;
let info = adapter.get_info();
eprintln!(
"[gpu] using {} ({:?}) via {:?}",
info.name, info.device_type, info.backend
);
let (device, queue) = adapter
.request_device(&wgpu::DeviceDescriptor {
label: Some(label),
required_features: wgpu::Features::empty(),
required_limits: wgpu::Limits::default(),
memory_hints: wgpu::MemoryHints::default(),
trace: wgpu::Trace::Off,
})
.await?;
Ok(GpuContext { device, queue })
}
/// Build a compute pipeline from an embedded WGSL source string.
pub fn compute_pipeline(
&self,
label: &str,
source: &'static str,
layout: &wgpu::PipelineLayout,
) -> wgpu::ComputePipeline {
let module = self
.device
.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some(label),
source: wgpu::ShaderSource::Wgsl(source.into()),
});
self.device
.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some(label),
layout: Some(layout),
module: &module,
entry_point: Some("main"),
compilation_options: Default::default(),
cache: None,
})
}
/// Create a storage buffer of `size` bytes (used for readback or GPU output).
pub fn create_storage_buffer(
&self,
label: &str,
size: u64,
contents: Option<&[u8]>,
) -> wgpu::Buffer {
let usage = wgpu::BufferUsages::STORAGE
| wgpu::BufferUsages::COPY_SRC
| wgpu::BufferUsages::COPY_DST;
match contents {
Some(data) => self
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some(label),
contents: data,
usage,
}),
None => self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some(label),
size,
usage,
mapped_at_creation: false,
}),
}
}
/// Create a uniform buffer initialised from `data`.
pub fn create_uniform_buffer(&self, label: &str, data: &[u8]) -> wgpu::Buffer {
let usage = wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST;
self.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some(label),
contents: data,
usage,
})
}
/// Create a read-only RGBA8 input texture (sampled by the compute kernels).
pub fn create_texture_2d(&self, label: &str, w: u32, h: u32) -> wgpu::Texture {
self.device.create_texture(&wgpu::TextureDescriptor {
label: Some(label),
size: wgpu::Extent3d {
width: w,
height: h,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::Rgba8Unorm,
usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
view_formats: &[],
})
}
/// Create a non-square output storage texture (written by kernels, read back).
pub fn create_storage_texture_2d_wh(&self, label: &str, w: u32, h: u32) -> wgpu::Texture {
self.device.create_texture(&wgpu::TextureDescriptor {
label: Some(label),
size: wgpu::Extent3d {
width: w,
height: h,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::Rgba8Unorm,
usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::COPY_SRC,
view_formats: &[],
})
}
/// Create an output storage texture (written by the kernels, read back).
pub fn create_storage_texture_2d(&self, label: &str, n: u32) -> wgpu::Texture {
self.device.create_texture(&wgpu::TextureDescriptor {
label: Some(label),
size: wgpu::Extent3d {
width: n,
height: n,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::Rgba8Unorm,
usage: wgpu::TextureUsages::STORAGE_BINDING | wgpu::TextureUsages::COPY_SRC,
view_formats: &[],
})
}
/// Upload RGBA8 `data` (width*height*4 bytes) into `tex`.
pub fn upload_rgba(&self, tex: &wgpu::Texture, w: u32, h: u32, data: &[u8]) {
self.queue.write_texture(
wgpu::TexelCopyTextureInfo {
texture: tex,
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
data,
wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(w * 4),
rows_per_image: Some(h),
},
wgpu::Extent3d {
width: w,
height: h,
depth_or_array_layers: 1,
},
);
}
/// Copy an RGBA storage texture to a tightly packed `Vec<u8>` (w*h*4).
///
/// Handles the WebGPU 256-byte `bytes_per_row` padding: a padded staging
/// buffer is mapped and then destrided into the tight layout expected by
/// the `image` crate.
pub fn read_texture_rgba(
&self,
tex: &wgpu::Texture,
w: u32,
h: u32,
) -> anyhow::Result<Vec<u8>> {
let bytes_per_row = u64::from(w) * 4;
let padded_row = align_bytes_per_row(bytes_per_row);
let total = padded_row * u64::from(h);
let staging = self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("texture-readback-staging"),
size: total,
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
mapped_at_creation: false,
});
let mut encoder = self
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
encoder.copy_texture_to_buffer(
wgpu::TexelCopyTextureInfo {
texture: tex,
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
wgpu::TexelCopyBufferInfo {
buffer: &staging,
layout: wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(padded_row as u32),
rows_per_image: Some(h),
},
},
wgpu::Extent3d {
width: w,
height: h,
depth_or_array_layers: 1,
},
);
self.queue.submit(Some(encoder.finish()));
let slice = staging.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |res| {
let _ = tx.send(res);
});
let _ = self.device.poll(wgpu::PollType::Wait);
rx.recv()??;
let mapped = slice.get_mapped_range();
let mut tight = vec![0u8; (bytes_per_row * u64::from(h)) as usize];
if padded_row == bytes_per_row {
tight.copy_from_slice(&mapped);
} else {
let row_w = bytes_per_row as usize;
let row_p = padded_row as usize;
for r in 0..h as usize {
let src_start = r * row_p;
tight[r * row_w..(r + 1) * row_w]
.copy_from_slice(&mapped[src_start..src_start + row_w]);
}
}
Ok(tight)
}
/// Map several already-submitted buffers in one `poll(Wait)` and return their
/// contents in order. Amortises the per-readback poll latency across a batch.
#[allow(clippy::needless_range_loop)]
pub fn readback_many(
&self,
srcs: &[&wgpu::Buffer],
sizes: &[u64],
) -> anyhow::Result<Vec<Vec<u8>>> {
assert_eq!(srcs.len(), sizes.len(), "readback_many: len mismatch");
let n = srcs.len();
if n == 0 {
return Ok(Vec::new());
}
let stagers: Vec<wgpu::Buffer> = (0..n)
.map(|i| {
self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("rb-many-staging"),
size: sizes[i],
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
mapped_at_creation: false,
})
})
.collect();
let mut encoder = self
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
for i in 0..n {
encoder.copy_buffer_to_buffer(srcs[i], 0, &stagers[i], 0, sizes[i]);
}
self.queue.submit(Some(encoder.finish()));
let (tx, rx) = std::sync::mpsc::channel();
for i in 0..n {
let tx = tx.clone();
stagers[i]
.slice(..)
.map_async(wgpu::MapMode::Read, move |res| {
let _ = tx.send((i, res));
});
}
let _ = self.device.poll(wgpu::PollType::Wait);
let mut results: Vec<Vec<u8>> = (0..n).map(|_| Vec::new()).collect();
for _ in 0..n {
let (i, res) = rx
.recv()
.map_err(|_| anyhow::anyhow!("readback_many: channel closed"))?;
res?;
results[i] = stagers[i].slice(..).get_mapped_range().to_vec();
}
Ok(results)
}
/// Copy several storage textures to tightly packed RGBA `Vec<u8>`s with one
/// submit + one poll. Handles the 256-byte `bytes_per_row` padding per tex.
#[allow(clippy::needless_range_loop)]
pub fn read_textures_rgba(
&self,
texs: &[&wgpu::Texture],
w: u32,
h: u32,
) -> anyhow::Result<Vec<Vec<u8>>> {
let n = texs.len();
if n == 0 {
return Ok(Vec::new());
}
let bytes_per_row = u64::from(w) * 4;
let padded_row = align_bytes_per_row(bytes_per_row);
let total = padded_row * u64::from(h);
let stagers: Vec<wgpu::Buffer> = (0..n)
.map(|_| {
self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("tex-rb-many-staging"),
size: total,
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
mapped_at_creation: false,
})
})
.collect();
let mut encoder = self
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
for i in 0..n {
encoder.copy_texture_to_buffer(
wgpu::TexelCopyTextureInfo {
texture: texs[i],
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
wgpu::TexelCopyBufferInfo {
buffer: &stagers[i],
layout: wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(padded_row as u32),
rows_per_image: Some(h),
},
},
wgpu::Extent3d {
width: w,
height: h,
depth_or_array_layers: 1,
},
);
}
self.queue.submit(Some(encoder.finish()));
let (tx, rx) = std::sync::mpsc::channel();
for i in 0..n {
let tx = tx.clone();
stagers[i]
.slice(..)
.map_async(wgpu::MapMode::Read, move |res| {
let _ = tx.send((i, res));
});
}
let _ = self.device.poll(wgpu::PollType::Wait);
let mut results: Vec<Vec<u8>> = (0..n).map(|_| Vec::new()).collect();
for _ in 0..n {
let (i, res) = rx
.recv()
.map_err(|_| anyhow::anyhow!("read_textures_rgba: channel closed"))?;
res?;
let mapped = stagers[i].slice(..).get_mapped_range();
let mut tight = vec![0u8; (bytes_per_row * u64::from(h)) as usize];
if padded_row == bytes_per_row {
tight.copy_from_slice(&mapped);
} else {
let row_w = bytes_per_row as usize;
let row_p = padded_row as usize;
for r in 0..h as usize {
let src_start = r * row_p;
tight[r * row_w..(r + 1) * row_w]
.copy_from_slice(&mapped[src_start..src_start + row_w]);
}
}
results[i] = tight;
}
Ok(results)
}
/// Snapshot the contents of a (storage or uniform) buffer into a `Vec<u8>`.
pub fn readback(&self, src: &wgpu::Buffer, size: u64) -> anyhow::Result<Vec<u8>> {
let snapshot = self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("readback"),
size,
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
mapped_at_creation: false,
});
let mut encoder = self
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
encoder.copy_buffer_to_buffer(src, 0, &snapshot, 0, size);
self.queue.submit(Some(encoder.finish()));
let slice = snapshot.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |res| {
let _ = tx.send(res);
});
let _ = self.device.poll(wgpu::PollType::Wait);
rx.recv()??;
let data = slice.get_mapped_range().to_vec();
Ok(data)
}
}
// Shader sources (embedded at compile time).
pub mod shaders {
/// Equirect -> 6 cube faces (images: bilinear; masks: nearest).
pub const EQUIRECT_TO_CUBEMAP: &str = include_str!("shaders/equirect_to_cubemap.wgsl");
/// Downscale + grayscale + Laplacian variance (sharpness scoring).
pub const SHARPNESS: &str = include_str!("shaders/sharpness.wgsl");
/// Downscale to a tiny luma thumbnail (motion/parallax proxy).
pub const THUMBNAIL: &str = include_str!("shaders/thumbnail.wgsl");
/// Cubemap -> equirect (nearest, for mask reprojection).
pub const C2E: &str = include_str!("shaders/c2e.wgsl");
}
+112
View File
@@ -0,0 +1,112 @@
// Stage 2 helper: cubemap -> equirect (NEAREST, for binary masks).
//
// The exact inverse of `equirect_to_cubemap.wgsl`'s `xyzcube` geometry, using a
// gnomonic (tangent) projection: the cube-surface coordinate for a face is
// 0.5 * (in-plane component / face-normal component), which inverts the linear
// cube-surface grid used by e2c. For each equirect output pixel we compute the
// 3D direction from (lon, lat), pick the dominant-axis cube face (F/R/B/L/U/D),
// invert that face's `xyzcube` mapping to a face pixel (j, i), and nearest-
// sample the matching face mask texture. This is the consistent inverse of our
// e2c kernel (round-trip e2c -> c2e reproduces the source within edge loss),
// and is close to py360convert.c2e(..., mode="nearest").
//
// Bindings 1..6 are the six face mask textures in FACES order [F,R,B,L,U,D],
// each `face_n x face_n`, RGBA8 (mask value replicated across RGB).
struct Params {
out_w: u32,
out_h: u32,
face_n: u32,
pad: u32,
}
@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var faceF: texture_2d<f32>;
@group(0) @binding(2) var faceR: texture_2d<f32>;
@group(0) @binding(3) var faceB: texture_2d<f32>;
@group(0) @binding(4) var faceL: texture_2d<f32>;
@group(0) @binding(5) var faceU: texture_2d<f32>;
@group(0) @binding(6) var faceD: texture_2d<f32>;
@group(0) @binding(7) var out_tex: texture_storage_2d<rgba8unorm, write>;
const PI: f32 = 3.141592653589793;
fn sample_face(id: u32, j: i32, i: i32) -> vec4<f32> {
let n = i32(params.face_n);
let cj = clamp(j, 0, n - 1);
let ci = clamp(i, 0, n - 1);
switch id {
case 0u: { return textureLoad(faceF, vec2<i32>(cj, ci), 0); }
case 1u: { return textureLoad(faceR, vec2<i32>(cj, ci), 0); }
case 2u: { return textureLoad(faceB, vec2<i32>(cj, ci), 0); }
case 3u: { return textureLoad(faceL, vec2<i32>(cj, ci), 0); }
case 4u: { return textureLoad(faceU, vec2<i32>(cj, ci), 0); }
case 5u: { return textureLoad(faceD, vec2<i32>(cj, ci), 0); }
default: { return vec4<f32>(0.0); }
}
}
@compute @workgroup_size(8, 8, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let w = params.out_w;
let h = params.out_h;
if (gid.x >= w || gid.y >= h) { return; }
let px = f32(gid.x);
let py = f32(gid.y);
let u_ = ((px + 0.5) / f32(w) - 0.5) * 2.0 * PI; // lon
let v_ = -((py + 0.5) / f32(h) - 0.5) * PI; // lat
let cv = cos(v_);
let dx = cv * sin(u_);
let dy = sin(v_);
let dz = cv * cos(u_);
let ax = abs(dx);
let ay = abs(dy);
let az = abs(dz);
let nm1 = f32(params.face_n - 1u);
var face_id: u32;
var jf: f32;
var if_: f32;
if (ay >= ax && ay >= az) {
if (dy > 0.0) {
// U: j=0.5*(x/y+1)*nm1, i=0.5*(z/y+1)*nm1
face_id = 4u;
jf = 0.5 * (dx / dy + 1.0) * nm1;
if_ = 0.5 * (dz / dy + 1.0) * nm1;
} else {
// D: j=0.5*(1-x/y)*nm1, i=0.5*(1+z/y)*nm1
face_id = 5u;
jf = 0.5 * (1.0 - dx / dy) * nm1;
if_ = 0.5 * (1.0 + dz / dy) * nm1;
}
} else if (ax >= az) {
if (dx > 0.0) {
// R: j=0.5*(1-z/x)*nm1, i=0.5*(1-y/x)*nm1
face_id = 1u;
jf = 0.5 * (1.0 - dz / dx) * nm1;
if_ = 0.5 * (1.0 - dy / dx) * nm1;
} else {
// L: j=0.5*(1-z/x)*nm1, i=0.5*(1+y/x)*nm1
face_id = 3u;
jf = 0.5 * (1.0 - dz / dx) * nm1;
if_ = 0.5 * (1.0 + dy / dx) * nm1;
}
} else {
if (dz > 0.0) {
// F: j=0.5*(x/z+1)*nm1, i=0.5*(1-y/z)*nm1
face_id = 0u;
jf = 0.5 * (dx / dz + 1.0) * nm1;
if_ = 0.5 * (1.0 - dy / dz) * nm1;
} else {
// B: j=0.5*(x/z+1)*nm1, i=0.5*(1+y/z)*nm1
face_id = 2u;
jf = 0.5 * (dx / dz + 1.0) * nm1;
if_ = 0.5 * (1.0 + dy / dz) * nm1;
}
}
let color = sample_face(face_id, i32(round(jf)), i32(round(if_)));
textureStore(out_tex, vec2<i32>(gid.xy), vec4<f32>(color.rgb, 1.0));
}
+145
View File
@@ -0,0 +1,145 @@
// Stage 3: project an equirectangular texture onto the 6 cube faces.
//
// This is a faithful GPU port of `py360convert.e2c` (mode "bilinear" for
// images, "nearest" for masks), verified against the installed
// py360convert source. The mapping is:
//
// 1. cube surface coords from `xyzcube(face_w)`:
// rng = linspace(-0.5, 0.5, num=N) (endpoint-INCLUSIVE)
// x_ = rng[j], y_ = -rng[i] = 0.5 - i/(N-1) (row 0 == top)
// per face (F R B L U D):
// F: (x_, y_, 0.5) R: ( 0.5, y_, -x_)
// B: (-x_, y_, -0.5) L: (-0.5, y_, x_)
// U: (x_, 0.5, y_) D: (x_, -0.5, -y_)
// 2. xyz2uv: u = atan2(x, z); c = hypot(x, z); v = atan2(y, c)
// 3. uv2coor: coor_x = (u/2pi + 0.5)*w - 0.5
// coor_y = (-v/pi + 0.5)*h - 0.5
// 4. sample a 1-px-PADDED version of the equirect with bilinear/nearest:
// - top/bottom pad rows are the image top/bottom row rotated by w/2
// (so the +180° seam is continuous),
// - left/right pad columns clamp to the adjacent interior column.
// The sampler's coordinate shift of +1 compensates for the padding.
//
// Storage textures are RGBA; for images the BGR<->RGB swap (OpenCV vs PNG)
// is done on the CPU after readback. Masks are returned as a single binary
// (0/255) value in all three channels so the CPU can take channel 0.
struct Params {
src_w: u32,
src_h: u32,
face_n: u32,
face_id: u32, // 0..5 in FACES order [F,R,B,L,U,D]
mode: u32, // 0 = bilinear (images), 1 = nearest (masks)
pad0: u32,
pad1: u32,
pad2: u32,
}
const PI: f32 = 3.141592653589793;
@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var src_tex: texture_2d<f32>;
@group(0) @binding(2) var src_samp: sampler;
@group(0) @binding(3) var out_tex: texture_storage_2d<rgba8unorm, write>;
// Sample the padded equirect. `i`/`j` are *padded-image* indices in
// [0, h+1] x [0, w+1]; the rules mirror py360convert's `EquirecSampler._pad`.
fn pad_sample(i: u32, j: u32) -> vec4<f32> {
let w = params.src_w;
let h = params.src_h;
let half = w / 2u;
if (i == 0u) {
// Top pad row = image row 0 rolled right by w/2.
let c = (j + half + w - 1u) % w;
return textureLoad(src_tex, vec2<i32>(i32(c), 0), 0);
}
if (i == h + 1u) {
let c = (j + half + w - 1u) % w;
return textureLoad(src_tex, vec2<i32>(i32(c), i32(h - 1u)), 0);
}
// Interior rows: i0 = image row.
let i0 = i - 1u;
if (j == 0u) {
// Left pad col = rightmost interior col.
return textureLoad(src_tex, vec2<i32>(i32(w - 1u), i32(i0)), 0);
}
if (j == w + 1u) {
return textureLoad(src_tex, vec2<i32>(0, i32(i0)), 0);
}
return textureLoad(src_tex, vec2<i32>(i32(j - 1u), i32(i0)), 0);
}
// Bilinear over the padded image. fx, fy are padded coords = coor + 1.
fn bilinear_pad(fx: f32, fy: f32) -> vec4<f32> {
let x0 = floor(fx);
let fracx = fx - x0;
let y0 = floor(fy);
let fracy = fy - y0;
let x1 = x0 + 1.0;
let y1 = y0 + 1.0;
let c00 = pad_sample(u32(y0), u32(x0));
let c10 = pad_sample(u32(y0), u32(x1));
let c01 = pad_sample(u32(y1), u32(x0));
let c11 = pad_sample(u32(y1), u32(x1));
let top = mix(c00, c10, fracx);
let bot = mix(c01, c11, fracx);
return mix(top, bot, fracy);
}
// Nearest over the padded image (cv2 INTER_NEAREST rounds to nearest pixel).
fn nearest_pad(fx: f32, fy: f32) -> vec4<f32> {
let j = i32(round(fx));
let i = i32(round(fy));
return pad_sample(u32(i), u32(j));
}
fn cubemap_color(i: u32, j: u32, n: u32) -> vec4<f32> {
// Endpoint-inclusive grid, matching xyzcube's linspace(-0.5, 0.5, N).
// rng_j is the column (image-right) coordinate; yv is the row coordinate.
// yv[i] = -rng[i], and rng_i is the reversed-row value used by U/D.
let rng_j = -0.5 + f32(j) / (f32(n) - 1.0);
let rng_i = -0.5 + f32(i) / (f32(n) - 1.0);
let yv = 0.5 - f32(i) / (f32(n) - 1.0);
var dir: vec3<f32>;
switch params.face_id {
case 0u: { dir = vec3<f32>(rng_j, yv, 0.5); } // F
case 1u: { dir = vec3<f32>(0.5, yv, -rng_j); } // R
case 2u: { dir = vec3<f32>(-rng_j, yv, -0.5); } // B
case 3u: { dir = vec3<f32>(-0.5, yv, rng_j); } // L
case 4u: { dir = vec3<f32>(rng_j, 0.5, rng_i); } // U (Z = fliplr(y) = rng_i)
case 5u: { dir = vec3<f32>(rng_j, -0.5, yv); } // D (Z = y = yv)
default: { dir = vec3<f32>(rng_j, yv, 0.5); }
}
// xyz2uv.
let u_ = atan2(dir.x, dir.z);
let c = sqrt(dir.x * dir.x + dir.z * dir.z);
let v_ = atan2(dir.y, c);
// uv2coor (original-image coordinates), then +1 for padding shift.
let wf = f32(params.src_w);
let hf = f32(params.src_h);
let coor_x = (u_ / (2.0 * PI) + 0.5) * wf - 0.5;
let coor_y = (-v_ / PI + 0.5) * hf - 0.5;
let fx = coor_x + 1.0;
let fy = coor_y + 1.0;
if (params.mode == 0u) {
return bilinear_pad(fx, fy);
} else {
return nearest_pad(fx, fy);
}
}
@compute @workgroup_size(8, 8, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let n = params.face_n;
if (gid.x >= n || gid.y >= n) {
return;
}
let color = cubemap_color(gid.y, gid.x, n);
// Opaque output; the BGR<->RGB swap happens on the CPU side.
textureStore(out_tex, vec2<i32>(gid.xy), vec4<f32>(color.rgb, 1.0));
}
+102
View File
@@ -0,0 +1,102 @@
// Stage 1 analysis kernel: per-frame quality metrics in one dispatch.
//
// For each 16x16 workgroup tile we compute, per pixel:
// g = luminance (0.299R + 0.587G + 0.114B)
// L = 3x3 4-neighbour Laplacian (matches cv2.Laplacian default ksize=1)
// and reduce within the workgroup (shared memory) to five partial sums:
// sum_g, sum_g2, sum_L, sum_L2, count
// One thread per workgroup writes these 5 floats to `partial[wg_index*5]`.
// The CPU sums the (small) partial buffer and derives:
// exposure = mean(g) = sum_g / n
// contrast = std(g) = sqrt(max(sum_g2/n - mean^2, 0))
// sharpness = var(L) = max(sum_L2/n - mean_L^2, 0)
// (Variance-of-Laplacian, the original Python metric.)
//
// This keeps the GPU->CPU readback tiny (5 floats/workgroup) instead of the
// full per-pixel Laplacian plane.
struct Params {
width: u32,
height: u32,
pad0: u32,
pad1: u32,
}
@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var src_tex: texture_2d<f32>;
@group(0) @binding(2) var<storage, read_write> partial: array<f32>;
const WG: u32 = 16u;
const WG2: u32 = WG * WG; // 256
var<workgroup> s_g: array<f32, WG2>;
var<workgroup> s_g2: array<f32, WG2>;
var<workgroup> s_l: array<f32, WG2>;
var<workgroup> s_l2: array<f32, WG2>;
var<workgroup> s_n: array<f32, WG2>;
fn luma(c: vec4<f32>) -> f32 {
return 0.299 * c.r + 0.587 * c.g + 0.114 * c.b;
}
fn load_clamp(x: i32, y: i32) -> vec4<f32> {
let cx = clamp(x, 0, i32(params.width) - 1);
let cy = clamp(y, 0, i32(params.height) - 1);
return textureLoad(src_tex, vec2<i32>(cx, cy), 0);
}
fn laplacian(x: i32, y: i32) -> f32 {
let c = luma(load_clamp(x, y));
let l = luma(load_clamp(x - 1, y));
let r = luma(load_clamp(x + 1, y));
let u = luma(load_clamp(x, y - 1));
let d = luma(load_clamp(x, y + 1));
return -4.0 * c + l + r + u + d;
}
@compute @workgroup_size(16, 16, 1)
fn main(
@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_index) lid: u32,
@builtin(global_invocation_id) gid: vec3<u32>,
) {
let gx = i32(gid.x);
let gy = i32(gid.y);
let inside = gx < i32(params.width) && gy < i32(params.height);
var g: f32 = 0.0;
var lp: f32 = 0.0;
if (inside) {
g = luma(load_clamp(gx, gy));
lp = laplacian(gx, gy);
}
s_g[lid] = g;
s_g2[lid] = g * g;
s_l[lid] = lp;
s_l2[lid] = lp * lp;
s_n[lid] = select(0.0, 1.0, inside);
workgroupBarrier();
// Tree reduction over 256 -> 1.
var stride = WG2 / 2u;
while (stride > 0u) {
if (lid < stride) {
s_g[lid] += s_g[lid + stride];
s_g2[lid] += s_g2[lid + stride];
s_l[lid] += s_l[lid + stride];
s_l2[lid] += s_l2[lid + stride];
s_n[lid] += s_n[lid + stride];
}
workgroupBarrier();
stride /= 2u;
}
if (lid == 0u) {
let idx = (wid.y * ((params.width + WG - 1u) / WG) + wid.x) * 5u;
partial[idx + 0u] = s_g[0];
partial[idx + 1u] = s_g2[0];
partial[idx + 2u] = s_l[0];
partial[idx + 3u] = s_l2[0];
partial[idx + 4u] = s_n[0];
}
}
+36
View File
@@ -0,0 +1,36 @@
// Stage-1 companion kernel: build a tiny grayscale luma thumbnail on the GPU.
//
// One thread per thumbnail pixel. Each thread nearest-samples a source pixel
// and writes its luma (0..255) into `thumb[ty * tw + tx]`. This lets the motion
// (parallax) proxy be computed on the GPU instead of the CPU reading the full
// RGBA frame, so extraction leans on the GPU rather than the CPU.
struct Params {
w: u32,
h: u32,
tw: u32,
th: u32,
pad0: u32,
}
@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var src_tex: texture_2d<f32>;
@group(0) @binding(2) var<storage, read_write> thumb: array<f32>;
fn luma(c: vec4<f32>) -> f32 {
return 0.299 * c.r + 0.587 * c.g + 0.114 * c.b;
}
@compute @workgroup_size(256, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let i = gid.x;
if (i >= params.tw * params.th) {
return;
}
let tx = i % params.tw;
let ty = i / params.tw;
let sx = i32(tx * params.w / params.tw);
let sy = i32(ty * params.h / params.th);
// Scale to 0..255 to match the CPU thumbnail / motion-diff convention.
thumb[i] = luma(textureLoad(src_tex, vec2<i32>(sx, sy), 0)) * 255.0;
}
+890
View File
@@ -0,0 +1,890 @@
//! Native GUI for the 360° -> COLMAP cubemap rig pipeline.
//!
//! Built with **egui** (rendered through **wgpu** via eframe) so it is
//! cross-vendor / cross-platform (Vulkan/Metal/DX12) just like the pipeline.
//! The GUI drives the pipeline by shelling out to the `pipeline_360` CLI binary
//! with `--gui-log-output`, streaming its `[progress] ` / log line protocol —
//! the same protocol the Python GUI uses — so the heavy GPU work runs in the
//! (optionally release-built) CLI process and the GUI stays responsive.
//!
//! Build: cargo build --features gui
//! Run: ./target/debug/pipeline_360_gui
//!
//! The GUI looks for the `pipeline_360` binary beside itself, then on PATH.
//! Native GUI (egui + wgpu). Enabled by the `gui` feature.
#![cfg(feature = "gui")]
use std::collections::VecDeque;
use std::path::PathBuf;
use std::process::Stdio;
use std::sync::mpsc::{self, Receiver, Sender};
use std::time::Instant;
use eframe::egui;
const PROGRESS_PREFIX: &str = "[progress] ";
const LOG_CAP: usize = 4000;
/// Locate the `pipeline_360` CLI binary: beside the GUI exe, else on PATH.
fn resolve_cli_binary() -> Option<PathBuf> {
if let Ok(exe) = std::env::current_exe() {
if let Some(dir) = exe.parent() {
for name in ["pipeline_360", "pipeline_360.exe"] {
let p = dir.join(name);
if p.exists() {
return Some(p);
}
}
}
}
which("pipeline_360")
}
/// Minimal `which` (no extra dep): search PATH for an executable.
fn which(name: &str) -> Option<PathBuf> {
let path = std::env::var_os("PATH")?;
for dir in std::env::split_paths(&path) {
let full = dir.join(name);
if full.is_file() {
return Some(full);
}
}
None
}
enum Line {
Log(String),
Progress(String),
Done(i32),
}
struct RunHandle {
child: std::process::Child,
rx: Receiver<Line>,
started: Instant,
code: Option<i32>,
}
impl RunHandle {
fn spawn(mut cmd: std::process::Command) -> std::io::Result<RunHandle> {
let mut child = cmd.stdout(Stdio::piped()).stderr(Stdio::piped()).spawn()?;
let stdout = child.stdout.take().unwrap();
let stderr = child.stderr.take().unwrap();
let (tx, rx) = mpsc::channel::<Line>();
spawn_line_reader(stdout, tx.clone(), false);
spawn_line_reader(stderr, tx, true);
Ok(RunHandle {
child,
rx,
started: Instant::now(),
code: None,
})
}
fn poll(&mut self, log: &mut VecDeque<String>) {
while let Ok(line) = self.rx.try_recv() {
match line {
Line::Log(s) | Line::Progress(s) => {
log.push_back(s);
while log.len() > LOG_CAP {
log.pop_front();
}
}
Line::Done(c) => self.code = Some(c),
}
}
if self.code.is_none() {
if let Ok(Some(status)) = self.child.try_wait() {
self.code = Some(status.code().unwrap_or(-1));
}
}
}
fn kill(&mut self) {
let _ = self.child.kill();
let _ = self.child.wait();
self.code = Some(-1);
}
fn elapsed(&self) -> std::time::Duration {
self.started.elapsed()
}
}
fn spawn_line_reader<R: std::io::Read + Send + 'static>(
mut r: R,
tx: Sender<Line>,
is_stderr: bool,
) {
std::thread::spawn(move || {
let mut buf = [0u8; 8192];
let mut leftover = Vec::new();
loop {
let n = match r.read(&mut buf) {
Ok(0) => break,
Ok(n) => n,
Err(_) => break,
};
leftover.extend_from_slice(&buf[..n]);
while let Some(pos) = leftover.iter().position(|&b| b == b'\n') {
let line_bytes = leftover.drain(..=pos).collect::<Vec<u8>>();
let raw = String::from_utf8_lossy(&line_bytes);
let raw = raw.trim_end_matches(['\n', '\r']);
if let Some(rest) = raw.strip_prefix(PROGRESS_PREFIX) {
let _ = tx.send(Line::Progress(format!("[…] {rest}")));
} else if !raw.is_empty() {
let _ = tx.send(Line::Log(if is_stderr {
format!("[stderr] {raw}")
} else {
raw.to_string()
}));
}
}
}
let _ = tx.send(Line::Done(0));
});
}
// ---------------------------------------------------------------------------
// Form state — mirrors the CLI flags + defaults.
// ---------------------------------------------------------------------------
#[derive(Clone)]
struct ExtractForm {
chunk_size: u32,
scale_width: u32,
output_format: String,
jpeg_quality: u32,
min_sharpness: f64,
min_exposure: f64,
max_exposure: f64,
min_contrast: f64,
min_motion: f64,
max_motion: f64,
selection: String,
target_count: u32,
min_gap: u32,
bad_frame_action: String,
motion_mode: String,
batch: u32,
}
impl Default for ExtractForm {
fn default() -> Self {
Self {
chunk_size: 30,
scale_width: 1920,
output_format: "png".into(),
jpeg_quality: 95,
min_sharpness: 0.0,
min_exposure: 0.0,
max_exposure: 255.0,
min_contrast: 0.0,
min_motion: 0.0,
max_motion: 1.0e12,
selection: "chunk".into(),
target_count: 0,
min_gap: 0,
bad_frame_action: "keep".into(),
motion_mode: "frame_diff".into(),
batch: 8,
}
}
}
#[derive(Clone)]
struct MasksForm {
classes: String,
model: PathBuf,
device: String,
conf: f64,
imgsz: u32,
dilate: u32,
skip_faces: String,
batch: u32,
empty_if_none: bool,
iou: f64,
// MIGraphX compiled-model cache directory (optional override; defaults to
// ~/.cache/pipeline_360/migraphx when unset).
ep_cache: Option<PathBuf>,
}
impl Default for MasksForm {
fn default() -> Self {
Self {
classes: "person,car,bus,truck,bicycle,motorcycle".into(),
model: PathBuf::from("yolo11x-seg.onnx"),
device: "cpu".into(),
conf: 0.35,
imgsz: 1280,
dilate: 6,
skip_faces: String::new(),
batch: 8,
empty_if_none: false,
iou: 0.45,
ep_cache: None,
}
}
}
#[derive(Clone)]
struct CubemapForm {
face_size: u32,
skip_faces: String,
jpeg_quality: u32,
batch: u32,
}
impl Default for CubemapForm {
fn default() -> Self {
Self {
face_size: 1600,
skip_faces: "U,D".into(),
jpeg_quality: 95,
batch: 4,
}
}
}
#[derive(Clone)]
struct RigForm {
ref_face: String,
intrinsics_convention: String,
link_mode: String,
}
impl Default for RigForm {
fn default() -> Self {
Self {
ref_face: "F".into(),
intrinsics_convention: "linspace".into(),
link_mode: "symlink".into(),
}
}
}
// ---------------------------------------------------------------------------
// App
// ---------------------------------------------------------------------------
struct App {
project_dir: PathBuf,
videos: Vec<PathBuf>,
start_stage: String,
skip_masks: bool,
extract: ExtractForm,
masks: MasksForm,
cubemap: CubemapForm,
rig: RigForm,
cli_binary: Option<PathBuf>,
cli_binary_str: String,
run: Option<RunHandle>,
log: VecDeque<String>,
status: String,
last_progress: String,
}
impl Default for App {
fn default() -> Self {
let cli = resolve_cli_binary();
Self {
project_dir: PathBuf::new(),
videos: Vec::new(),
start_stage: "extract".into(),
skip_masks: false,
extract: ExtractForm::default(),
masks: MasksForm::default(),
cubemap: CubemapForm::default(),
rig: RigForm::default(),
cli_binary_str: cli
.as_ref()
.map(|p| p.display().to_string())
.unwrap_or_else(|| "pipeline_360 (NOT FOUND — build it first)".into()),
cli_binary: cli,
log: VecDeque::new(),
run: None,
status: "Idle".into(),
last_progress: String::new(),
}
}
}
impl App {
fn is_running(&self) -> bool {
self.run.as_ref().is_some_and(|h| h.code.is_none())
}
fn build_command(&self) -> Option<std::process::Command> {
let bin = self.cli_binary.clone()?;
let mut cmd = std::process::Command::new(bin);
cmd.arg("all");
cmd.arg("--project-dir").arg(&self.project_dir);
for v in &self.videos {
cmd.arg("--video").arg(v);
}
cmd.arg("--start-stage").arg(&self.start_stage);
if self.skip_masks {
cmd.arg("--skip-masks");
}
cmd.arg("--gui-log-output");
let e = &self.extract;
cmd.args([
"--extract-chunk-size",
&e.chunk_size.to_string(),
"--extract-scale-width",
&e.scale_width.to_string(),
"--extract-output-format",
&e.output_format,
"--extract-jpeg-quality",
&e.jpeg_quality.to_string(),
"--extract-min-sharpness",
&fmt_f(e.min_sharpness),
"--extract-min-exposure",
&fmt_f(e.min_exposure),
"--extract-max-exposure",
&fmt_f(e.max_exposure),
"--extract-min-contrast",
&fmt_f(e.min_contrast),
"--extract-min-motion",
&fmt_f(e.min_motion),
"--extract-max-motion",
&fmt_f(e.max_motion),
"--extract-selection",
&e.selection,
"--extract-target-count",
&e.target_count.to_string(),
"--extract-min-gap",
&e.min_gap.to_string(),
"--extract-bad-frame-action",
&e.bad_frame_action,
"--extract-motion-mode",
&e.motion_mode,
"--extract-batch",
&e.batch.to_string(),
]);
let m = &self.masks;
cmd.args([
"--mask-classes",
&m.classes,
"--mask-model",
&m.model.display().to_string(),
"--mask-device",
&m.device,
"--mask-conf",
&fmt_f(m.conf),
"--mask-imgsz",
&m.imgsz.to_string(),
"--mask-dilate",
&m.dilate.to_string(),
"--mask-skip-faces",
&m.skip_faces,
"--mask-batch",
&m.batch.to_string(),
"--mask-iou",
&fmt_f(m.iou),
]);
if m.empty_if_none {
cmd.arg("--mask-empty-if-none");
}
if let Some(cache) = &m.ep_cache {
cmd.arg("--mask-ep-cache").arg(cache);
}
let c = &self.cubemap;
cmd.args([
"--cubemap-face-size",
&c.face_size.to_string(),
"--cubemap-skip-faces",
&c.skip_faces,
"--cubemap-jpeg-quality",
&c.jpeg_quality.to_string(),
"--cubemap-batch",
&c.batch.to_string(),
]);
let r = &self.rig;
cmd.args([
"--rig-ref-face",
&r.ref_face,
"--rig-intrinsics-convention",
&r.intrinsics_convention,
"--rig-link-mode",
&r.link_mode,
]);
Some(cmd)
}
fn start(&mut self) {
if self.is_running() {
return;
}
if self.cli_binary.is_none() {
self.status = "Error: pipeline_360 binary not found (build it first).".into();
return;
}
if self.project_dir.as_os_str().is_empty() || self.videos.is_empty() {
self.status = "Error: choose a project dir and at least one video.".into();
return;
}
let Some(cmd) = self.build_command() else {
return;
};
self.log.clear();
self.last_progress.clear();
self.status = "Running…".into();
self.log.push_back(format!("$ {}", cmd_to_string(&cmd)));
match RunHandle::spawn(cmd) {
Ok(h) => self.run = Some(h),
Err(e) => self.status = format!("Error: failed to start pipeline: {e}"),
}
}
fn cancel(&mut self) {
if let Some(h) = self.run.as_mut() {
h.kill();
self.status = "Cancelled.".into();
}
}
}
fn fmt_f(v: impl ToString) -> String {
let s = v.to_string();
if s.eq_ignore_ascii_case("inf") {
"1e12".into()
} else {
s
}
}
fn cmd_to_string(cmd: &std::process::Command) -> String {
let prog = cmd.get_program().to_string_lossy().to_string();
let args: Vec<String> = cmd
.get_args()
.map(|a| {
let a = a.to_string_lossy().to_string();
if a.contains(' ') {
format!("\"{a}\"")
} else {
a
}
})
.collect();
format!("{prog} {}", args.join(" "))
}
fn combo(ui: &mut egui::Ui, id: &str, options: &[&str], val: &mut String) {
egui::ComboBox::from_id_salt(id)
.selected_text(val.clone())
.show_ui(ui, |ui| {
for o in options {
ui.selectable_value(val, (*o).to_string(), *o);
}
});
}
impl eframe::App for App {
fn ui(&mut self, ui: &mut egui::Ui, _frame: &mut eframe::Frame) {
// Poll subprocess each frame.
if let Some(h) = self.run.as_mut() {
h.poll(&mut self.log);
if let Some(code) = h.code {
if !self.status.starts_with("Cancelled") {
self.status = if code == 0 {
"Pipeline complete.".into()
} else {
format!("Pipeline exited with code {code}. See log.")
};
}
}
}
// Bottom: log panel + run controls.
egui::Panel::bottom("log_panel")
.resizable(true)
.min_size(150.0)
.show(ui, |ui| {
ui.horizontal(|ui| {
let status = if self.is_running() {
if let Some(h) = self.run.as_ref() {
format!("{} ({}s)", self.status, h.elapsed().as_secs())
} else {
self.status.clone()
}
} else {
self.status.clone()
};
ui.label(egui::RichText::new(status).strong());
ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
let running = self.is_running();
if running {
ui.spinner();
}
if ui
.add_enabled(running, egui::Button::new("⏹ Stop"))
.clicked()
{
self.cancel();
}
if ui
.add_enabled(!running, egui::Button::new("▶ Run pipeline"))
.clicked()
{
self.start();
}
});
});
ui.separator();
egui::ScrollArea::vertical()
.auto_shrink([false, false])
.stick_to_bottom(true)
.show(ui, |ui| {
ui.style_mut().wrap_mode = Some(egui::TextWrapMode::Wrap);
for line in &self.log {
if line.starts_with("[…] ") {
continue;
}
ui.label(egui::RichText::new(line).monospace().size(12.0));
}
if !self.last_progress.is_empty() {
ui.label(
egui::RichText::new(&self.last_progress)
.monospace()
.color(egui::Color32::LIGHT_BLUE)
.size(12.0),
);
}
});
});
// Left: inputs.
egui::Panel::left("inputs")
.resizable(true)
.min_size(320.0)
.show(ui, |ui| {
egui::ScrollArea::vertical().show(ui, |ui| {
ui.heading("360° -> COLMAP cubemap rig");
ui.add_space(2.0);
ui.label(
egui::RichText::new(format!("CLI: {}", self.cli_binary_str))
.small()
.color(egui::Color32::GRAY),
);
ui.add_space(8.0);
ui.label("Project directory");
ui.horizontal(|ui| {
let mut s = self.project_dir.display().to_string();
let r = ui.text_edit_singleline(&mut s);
if r.changed() {
self.project_dir = PathBuf::from(s);
}
if ui.button("Browse…").clicked() {
if let Some(p) = rfd::FileDialog::new().pick_folder() {
self.project_dir = p;
}
}
});
ui.add_space(8.0);
ui.label("Source videos");
let vids = self.videos.clone();
for (i, v) in vids.iter().enumerate() {
ui.horizontal(|ui| {
ui.label(format!("{}. {}", i + 1, v.display()));
if ui.small_button("✕").clicked() {
self.videos.remove(i);
}
});
}
if ui.button("+ Add videos…").clicked() {
if let Some(files) = rfd::FileDialog::new()
.add_filter("Videos", &["mp4", "mov", "mkv", "avi", "webm", "m4v"])
.pick_files()
{
self.videos.extend(files);
}
}
ui.add_space(8.0);
ui.horizontal(|ui| {
ui.label("Start stage:");
egui::ComboBox::from_id_salt("start_stage")
.selected_text(self.start_stage.clone())
.show_ui(ui, |ui| {
for s in ["extract", "cubemap", "masks", "rig"] {
ui.selectable_value(&mut self.start_stage, s.into(), s);
}
});
});
ui.checkbox(&mut self.skip_masks, "Skip masks stage");
});
});
// Centre: stage option sections.
egui::CentralPanel::default().show(ui, |ui| {
egui::ScrollArea::vertical().show(ui, |ui| {
ui.spacing_mut().item_spacing.y = 6.0;
egui::CollapsingHeader::new("Stage 1 — Extract (sharpest/best frames)")
.default_open(true)
.show(ui, |ui| {
let e = &mut self.extract;
egui::Grid::new("extract_grid")
.num_columns(2)
.spacing([8.0, 4.0])
.show(ui, |ui| {
ui.label("Chunk size");
ui.add(egui::DragValue::new(&mut e.chunk_size).range(1..=10000));
ui.end_row();
ui.label("Scale width");
ui.add(egui::DragValue::new(&mut e.scale_width).range(0..=8000));
ui.end_row();
ui.label("Output format");
combo(ui, "extract_fmt", &["png", "jpg"], &mut e.output_format);
ui.end_row();
ui.label("JPEG quality");
ui.add(egui::DragValue::new(&mut e.jpeg_quality).range(1..=100));
ui.end_row();
ui.label("Batch (frames/GPU)");
ui.add(egui::DragValue::new(&mut e.batch).range(1..=64));
ui.end_row();
ui.label("Selection");
combo(ui, "extract_sel", &["chunk", "spread"], &mut e.selection);
ui.end_row();
ui.label("Target count");
ui.add(egui::DragValue::new(&mut e.target_count).range(0..=100000));
ui.end_row();
ui.label("Min gap (frames)");
ui.add(egui::DragValue::new(&mut e.min_gap).range(0..=10000));
ui.end_row();
ui.label("Bad-frame action");
combo(
ui,
"extract_bfa",
&["keep", "skip", "warn"],
&mut e.bad_frame_action,
);
ui.end_row();
ui.label("Motion mode");
combo(ui, "extract_mm", &["frame_diff", "off"], &mut e.motion_mode);
ui.end_row();
});
ui.label(
egui::RichText::new("Quality thresholds (0/255 bounds = no filter):")
.small(),
);
let e = &mut self.extract;
egui::Grid::new("extract_q")
.num_columns(2)
.spacing([8.0, 4.0])
.show(ui, |ui| {
ui.label("Min sharpness");
ui.add(
egui::DragValue::new(&mut e.min_sharpness)
.speed(0.05)
.range(0.0..=1.0e9),
);
ui.end_row();
ui.label("Min exposure");
ui.add(
egui::DragValue::new(&mut e.min_exposure)
.speed(1.0)
.range(0.0..=255.0),
);
ui.end_row();
ui.label("Max exposure");
ui.add(
egui::DragValue::new(&mut e.max_exposure)
.speed(1.0)
.range(0.0..=255.0),
);
ui.end_row();
ui.label("Min contrast");
ui.add(
egui::DragValue::new(&mut e.min_contrast)
.speed(1.0)
.range(0.0..=255.0),
);
ui.end_row();
ui.label("Min motion");
ui.add(
egui::DragValue::new(&mut e.min_motion)
.speed(1.0)
.range(0.0..=255.0),
);
ui.end_row();
ui.label("Max motion");
ui.add(
egui::DragValue::new(&mut e.max_motion)
.speed(1.0)
.range(0.0..=1.0e9),
);
ui.end_row();
});
});
egui::CollapsingHeader::new(
"Stage 3 — Masks (YOLO segmentation; runs on cube faces)",
)
.show(ui, |ui| {
let m = &mut self.masks;
egui::Grid::new("masks_grid")
.num_columns(2)
.spacing([8.0, 4.0])
.show(ui, |ui| {
ui.label("Classes");
ui.text_edit_singleline(&mut m.classes);
ui.end_row();
ui.label("Model (ONNX)");
ui.horizontal(|ui| {
let mut s = m.model.display().to_string();
let r = ui.text_edit_singleline(&mut s);
if r.changed() {
m.model = PathBuf::from(s);
}
if ui.button("…").clicked() {
if let Some(p) = rfd::FileDialog::new()
.add_filter("ONNX", &["onnx"])
.pick_file()
{
m.model = p;
}
}
});
ui.end_row();
ui.label("Device");
combo(ui, "mask_dev", &["cpu", "rocm", "migraphx"], &mut m.device);
ui.end_row();
ui.label("Conf");
ui.add(
egui::DragValue::new(&mut m.conf)
.speed(0.01)
.range(0.0..=1.0)
.fixed_decimals(2),
);
ui.end_row();
ui.label("Imgsz");
ui.add(egui::DragValue::new(&mut m.imgsz).range(64..=4096));
ui.end_row();
ui.label("Dilate (px)");
ui.add(egui::DragValue::new(&mut m.dilate).range(0..=100));
ui.end_row();
ui.label("Skip faces");
ui.text_edit_singleline(&mut m.skip_faces);
ui.end_row();
ui.label("Batch (faces/GPU)");
ui.add(egui::DragValue::new(&mut m.batch).range(1..=64));
ui.end_row();
ui.label("NMS IoU");
ui.add(
egui::DragValue::new(&mut m.iou)
.speed(0.01)
.range(0.0..=1.0)
.fixed_decimals(2),
);
ui.end_row();
});
// Optional override for the MIGraphX compiled-model cache dir.
// Leave empty to use the default (~/.cache/pipeline_360/migraphx).
let mut cache_s = m
.ep_cache
.as_ref()
.map(|p| p.display().to_string())
.unwrap_or_default();
ui.horizontal(|ui| {
ui.label("GPU model cache dir:");
let r = ui.text_edit_singleline(&mut cache_s);
if r.changed() {
m.ep_cache = if cache_s.trim().is_empty() {
None
} else {
Some(PathBuf::from(cache_s.trim()))
};
}
});
ui.checkbox(
&mut m.empty_if_none,
"Write empty (all-keep) mask when no objects detected",
);
});
egui::CollapsingHeader::new("Stage 2 — Cubemap (equirect -> 6 faces)").show(
ui,
|ui| {
let c = &mut self.cubemap;
egui::Grid::new("cubemap_grid")
.num_columns(2)
.spacing([8.0, 4.0])
.show(ui, |ui| {
ui.label("Face size");
ui.add(egui::DragValue::new(&mut c.face_size).range(64..=4096));
ui.end_row();
ui.label("Skip faces");
ui.text_edit_singleline(&mut c.skip_faces);
ui.end_row();
ui.label("JPEG quality");
ui.add(egui::DragValue::new(&mut c.jpeg_quality).range(1..=100));
ui.end_row();
ui.label("Batch (frames/GPU)");
ui.add(egui::DragValue::new(&mut c.batch).range(1..=32));
ui.end_row();
});
},
);
egui::CollapsingHeader::new("Stage 4 — Rig config").show(ui, |ui| {
let r = &mut self.rig;
egui::Grid::new("rig_grid")
.num_columns(2)
.spacing([8.0, 4.0])
.show(ui, |ui| {
ui.label("Reference face");
combo(
ui,
"rig_ref",
&["F", "R", "B", "L", "U", "D"],
&mut r.ref_face,
);
ui.end_row();
ui.label("Intrinsics");
combo(
ui,
"rig_int",
&["linspace", "half"],
&mut r.intrinsics_convention,
);
ui.end_row();
ui.label("Link mode");
combo(
ui,
"rig_link",
&["symlink", "hardlink", "copy"],
&mut r.link_mode,
);
ui.end_row();
});
});
});
});
// Keep repainting while a run is in progress so the log streams live.
if self.is_running() {
ui.ctx().request_repaint();
}
}
}
pub fn launch() -> eframe::Result<()> {
let options = eframe::NativeOptions {
viewport: egui::ViewportBuilder::default()
.with_inner_size([1100.0, 760.0])
.with_min_inner_size([700.0, 500.0]),
..Default::default()
};
eframe::run_native(
"pipeline_360 — 360° -> COLMAP cubemap rig",
options,
Box::new(|_cc| Ok(Box::<App>::default())),
)
}
+318
View File
@@ -0,0 +1,318 @@
//! Minimal pure-`std` JSON value + writer.
//!
//! The rig config schema is small and fixed, so we don't need a full JSON
//! library, and keeping this pure-`std` lets the whole stage-4 core compile and
//! test offline. Used by `stages::rig` (writer) and `selftest` (reader).
/// A JSON value, big enough for the COLMAP rig schema.
#[derive(Debug, Clone, PartialEq)]
pub enum Value {
Object(Vec<(String, Value)>),
Array(Vec<Value>),
Number(f64),
Bool(bool),
String(String),
Null,
}
impl Value {
/// Depth-first lookup of the first object value at `key` (recursively).
pub fn find(&self, key: &str) -> Option<&Value> {
match self {
Value::Object(pairs) => {
if let Some((_, v)) = pairs.iter().find(|(k, _)| k == key) {
return Some(v);
}
for (_, v) in pairs {
if let Some(hit) = v.find(key) {
return Some(hit);
}
}
None
}
Value::Array(items) => items.iter().find_map(|v| v.find(key)),
_ => None,
}
}
pub fn as_array(&self) -> Option<&[Value]> {
match self {
Value::Array(items) => Some(items),
_ => None,
}
}
pub fn as_str(&self) -> Option<&str> {
match self {
Value::String(s) => Some(s),
_ => None,
}
}
pub fn as_f64(&self) -> Option<f64> {
match self {
Value::Number(n) => Some(*n),
_ => None,
}
}
}
fn write_string(out: &mut String, s: &str) {
out.push('"');
for c in s.chars() {
match c {
'"' => out.push_str("\\\""),
'\\' => out.push_str("\\\\"),
'\n' => out.push_str("\\n"),
'\r' => out.push_str("\\r"),
'\t' => out.push_str("\\t"),
c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
c => out.push(c),
}
}
out.push('"');
}
impl Value {
/// Compact JSON serialisation.
pub fn to_json(&self) -> String {
let mut out = String::new();
self.write(&mut out);
out
}
/// Serialise this value into `out` (used by the pretty printer too).
pub fn write(&self, out: &mut String) {
match self {
Value::Null => out.push_str("null"),
Value::Bool(b) => out.push_str(if *b { "true" } else { "false" }),
Value::Number(n) => {
if n.fract() == 0.0 && n.abs() < 1e15 {
out.push_str(&format!("{}", *n as i64));
} else {
out.push_str(&format!("{}", n));
}
}
Value::String(s) => write_string(out, s),
Value::Array(items) => {
out.push('[');
for (i, v) in items.iter().enumerate() {
if i > 0 {
out.push(',');
}
v.write(out);
}
out.push(']');
}
Value::Object(pairs) => {
out.push('{');
for (i, (k, v)) in pairs.iter().enumerate() {
if i > 0 {
out.push(',');
}
write_string(out, k);
out.push(':');
v.write(out);
}
out.push('}');
}
}
}
}
// ---------------------------------------------------------------------------
// A tiny recursive-descent parser (only as much as the writer + rig tests need).
// ---------------------------------------------------------------------------
/// Parse a JSON string into a `Value`. Returns an error message on failure.
pub fn parse(input: &str) -> Result<Value, String> {
let bytes = input.as_bytes();
let mut pos = 0usize;
let value = parse_value(bytes, &mut pos)?;
skip_ws(bytes, &mut pos);
if pos != bytes.len() {
return Err(format!("trailing characters at offset {pos}"));
}
Ok(value)
}
fn skip_ws(b: &[u8], pos: &mut usize) {
while *pos < b.len() && matches!(b[*pos], b' ' | b'\t' | b'\n' | b'\r') {
*pos += 1;
}
}
fn parse_value(b: &[u8], pos: &mut usize) -> Result<Value, String> {
skip_ws(b, pos);
if *pos >= b.len() {
return Err("unexpected end of input".into());
}
match b[*pos] {
b'{' => parse_object(b, pos),
b'[' => parse_array(b, pos),
b'"' => Ok(Value::String(parse_string(b, pos)?)),
b't' => {
expect(b, pos, "true")?;
Ok(Value::Bool(true))
}
b'f' => {
expect(b, pos, "false")?;
Ok(Value::Bool(false))
}
b'n' => {
expect(b, pos, "null")?;
Ok(Value::Null)
}
_ => parse_number(b, pos),
}
}
fn expect(b: &[u8], pos: &mut usize, word: &str) -> Result<(), String> {
if b[*pos..].starts_with(word.as_bytes()) {
*pos += word.len();
Ok(())
} else {
Err(format!("expected {word} at offset {pos}"))
}
}
fn parse_string(b: &[u8], pos: &mut usize) -> Result<String, String> {
if b[*pos] != b'"' {
return Err(format!("expected string at offset {pos}"));
}
*pos += 1;
let mut out = String::new();
while *pos < b.len() {
match b[*pos] {
b'"' => {
*pos += 1;
return Ok(out);
}
b'\\' => {
*pos += 1;
if *pos >= b.len() {
return Err("unterminated escape".into());
}
let c = b[*pos];
*pos += 1;
match c {
b'"' => out.push('"'),
b'\\' => out.push('\\'),
b'/' => out.push('/'),
b'n' => out.push('\n'),
b'r' => out.push('\r'),
b't' => out.push('\t'),
b'u' => {
if *pos + 4 > b.len() {
return Err("bad \\u escape".into());
}
let hex = std::str::from_utf8(&b[*pos..*pos + 4])
.map_err(|_| "bad utf8 in escape".to_string())?;
let code =
u32::from_str_radix(hex, 16).map_err(|_| "bad \\u hex".to_string())?;
*pos += 4;
out.push(char::from_u32(code).unwrap_or('\u{FFFD}'));
}
other => return Err(format!("bad escape \\{}", other as char)),
}
}
c => {
// Copy a UTF-8 sequence byte-by-byte (ASCII fast path).
if c < 0x80 {
out.push(c as char);
*pos += 1;
} else {
let rest = &b[*pos..];
let len = utf8_len(rest);
let s =
std::str::from_utf8(&rest[..len]).map_err(|_| "bad utf8".to_string())?;
out.push_str(s);
*pos += len;
}
}
}
}
Err("unterminated string".into())
}
fn utf8_len(b: &[u8]) -> usize {
if b.is_empty() {
return 0;
}
let lead = b[0];
let n = if lead >= 0xF0 {
4
} else if lead >= 0xE0 {
3
} else if lead >= 0xC0 {
2
} else {
1
};
n.min(b.len())
}
fn parse_number(b: &[u8], pos: &mut usize) -> Result<Value, String> {
let start = *pos;
while *pos < b.len() && matches!(b[*pos], b'0'..=b'9' | b'-' | b'+' | b'.' | b'e' | b'E') {
*pos += 1;
}
let s = std::str::from_utf8(&b[start..*pos]).map_err(|_| "bad number utf8".to_string())?;
let n: f64 = s.parse().map_err(|_| format!("bad number {s:?}"))?;
Ok(Value::Number(n))
}
fn parse_array(b: &[u8], pos: &mut usize) -> Result<Value, String> {
*pos += 1; // consume '['
let mut items = Vec::new();
loop {
skip_ws(b, pos);
if *pos >= b.len() {
return Err("unterminated array".into());
}
if b[*pos] == b']' {
*pos += 1;
return Ok(Value::Array(items));
}
if !items.is_empty() {
if b[*pos] != b',' {
return Err(format!("expected ',' in array at offset {pos}"));
}
*pos += 1;
}
items.push(parse_value(b, pos)?);
}
}
fn parse_object(b: &[u8], pos: &mut usize) -> Result<Value, String> {
*pos += 1; // consume '{'
let mut pairs = Vec::new();
loop {
skip_ws(b, pos);
if *pos >= b.len() {
return Err("unterminated object".into());
}
if b[*pos] == b'}' {
*pos += 1;
return Ok(Value::Object(pairs));
}
if !pairs.is_empty() {
if b[*pos] != b',' {
return Err(format!("expected ',' in object at offset {pos}"));
}
*pos += 1;
skip_ws(b, pos);
}
if b[*pos] != b'"' {
return Err(format!("expected key string at offset {pos}"));
}
let key = parse_string(b, pos)?;
skip_ws(b, pos);
if *pos >= b.len() || b[*pos] != b':' {
return Err(format!("expected ':' after key at offset {pos}"));
}
*pos += 1;
let value = parse_value(b, pos)?;
pairs.push((key, value));
}
}
+24
View File
@@ -0,0 +1,24 @@
//! `pipeline_360` — Rust/WebGPU rewrite.
//!
//! See `PLAN.md` for the full design. This crate mirrors the Python pipeline
//! (`pipeline_360.py`) stage-for-stage:
//!
//! 1. extract — sharpest frame per chunk (GPU variance-of-Laplacian)
//! 2. masks — YOLO segmentation (optional)
//! 3. cubemap — equirect → 6 cube faces (GPU projection)
//! 4. rig — COLMAP rig_config.json + nested `{FACE}/{stem}` tree
//!
//! The stage-4 geometry + rig generation (`math`, `json`, `stages::rig`) is
//! pure `std` and independently verifiable offline.
pub mod error;
#[cfg(feature = "gui")]
pub mod gui;
pub mod gpu;
pub mod json;
pub mod logging;
pub mod math;
pub mod project;
pub mod selftest;
pub mod stages;
pub mod video;
+33
View File
@@ -0,0 +1,33 @@
//! Logging and progress line protocol (pure `std`).
//!
//! The GUI (Python or future native) consumes plain stdout lines as log text
//! and lines prefixed with `PROGRESS_PREFIX` as in-place "replace the last
//! line" progress. Keep this byte-compatible with
//! `pipeline_360_gui.py`'s expectations.
/// Prefix marking a progress (replace-last-line) message.
pub const PROGRESS_PREFIX: &str = "[progress] ";
/// A sink for emitted messages. `None` means raw stdout.
pub type Logger<'a> = Option<&'a dyn Fn(&str)>;
/// Emit a normal log line (message-only; the CLI/GUI decides whether to add a newline).
pub fn emit(message: &str, logger: Logger) {
match logger {
Some(f) => f(message),
None => println!("{message}"),
}
}
/// Emit a progress (replace-last-line) message with the `[progress] ` prefix.
pub fn emit_progress(message: &str, logger: Logger) {
if let Some(f) = logger {
f(&format!("{PROGRESS_PREFIX}{message}"));
}
}
/// A stage header line, e.g. `==== STAGE 1/4: Extract Sharpest Frames ====`.
pub fn stage_header(index: usize, total: usize, title: &str, logger: Logger) {
emit("", logger);
emit(&format!("==== STAGE {index}/{total}: {title} ===="), logger);
}
+777
View File
@@ -0,0 +1,777 @@
//! Entry point: parse CLI, run the requested stage(s), stream progress.
use std::path::PathBuf;
use std::process::ExitCode;
use clap::{Args, Parser, Subcommand};
use pipeline_360::error::{CancelFlag, PipelineError};
use pipeline_360::gpu::GpuContext;
use pipeline_360::logging::{self, Logger};
use pipeline_360::project::ProjectLayout;
use pipeline_360::stages::cubemap;
use pipeline_360::stages::extract::{self, BadFrameAction, ExtractConfig, MotionMode, Selection};
use pipeline_360::stages::masks::{self, MasksConfig, Projection};
use pipeline_360::stages::rig::{self, Layout, LinkMode};
/// CLI parsed into a typed struct. `gui_log_output` is read from either the
/// parent or subcommand position.
#[derive(Debug, Parser)]
#[command(
name = "pipeline_360",
version,
about = "360° -> COLMAP cubemap rig pipeline"
)]
struct Cli {
/// Whether to emit `[progress] ` lines for the GUI (accepted in any position).
#[arg(long, hide = true)]
gui_log_output: bool,
/// Launch the native GUI (requires the `gui` feature). Also launched when
/// the binary is run with no subcommand.
#[arg(long, hide = true)]
gui: bool,
#[command(subcommand)]
command: Option<Command>,
}
#[derive(Debug, Subcommand)]
#[allow(clippy::large_enum_variant)]
enum Command {
/// Run the full pipeline: extract -> masks -> cubemap -> rig.
All(AllArgs),
/// Stage 1 only: extract sharpest/best frames from video.
Extract(ExtractArgs),
/// Stage 2 only: generate YOLO SfM masks (planned).
Masks(MasksArgs),
/// Stage 3 only: project equirect frames to cube faces.
Cubemap(CubemapArgs),
/// Stage 4 only: generate COLMAP rig config + nested tree (or --self-test).
Rig(RigArgs),
}
#[derive(Debug, Args, Clone)]
struct ExtractOpts {
#[arg(
long,
default_value_t = 30,
help = "Select one frame per this many frames (chunk mode)"
)]
chunk_size: u32,
#[arg(
long,
default_value_t = 1920,
help = "Resize width for analysis; 0 = native"
)]
scale_width: u32,
#[arg(
long,
default_value_t = 4,
help = "Worker count (parity; GPU is the parallelism)"
)]
workers: u32,
#[arg(
long,
default_value = "output_frame_%05d.png",
help = "printf-style output pattern"
)]
output_pattern: String,
#[arg(long, default_value = "png", help = "png | jpg")]
output_format: String,
#[arg(long, default_value_t = 95)]
jpeg_quality: u32,
#[arg(
long = "extract-regenerate-metadata",
help = "Ignore existing metadata CSV and regenerate"
)]
reuse_metadata_inv: bool,
// --- Quality thresholds ("throw away bad frames"). ---
#[arg(
long,
default_value_t = 0.0,
help = "Drop frames whose variance-of-Laplacian is below this"
)]
min_sharpness: f64,
#[arg(
long,
default_value_t = 0.0,
help = "Drop frames whose mean luma is below this (0..255)"
)]
min_exposure: f64,
#[arg(
long,
default_value_t = 255.0,
help = "Drop frames whose mean luma is above this (0..255)"
)]
max_exposure: f64,
#[arg(
long,
default_value_t = 0.0,
help = "Drop frames whose luma std is below this"
)]
min_contrast: f64,
#[arg(
long,
default_value_t = 0.0,
help = "Min inter-frame motion (0..255) — parallax baseline"
)]
min_motion: f64,
#[arg(
long,
default_value_t = 1.0e12,
help = "Max inter-frame motion (0..255) — reject jitter/blur"
)]
max_motion: f64,
// --- Selection ---
#[arg(long, default_value = "chunk", help = "chunk | spread")]
selection: String,
#[arg(
long,
default_value_t = 0,
help = "Target frame count (spread mode); 0 = all survivors / chunk-based"
)]
target_count: u32,
#[arg(
long,
default_value_t = 0,
help = "Min frame-index spacing between selected frames"
)]
min_gap: u32,
#[arg(
long,
default_value = "keep",
help = "skip | keep | warn — chunk with no surviving frame"
)]
bad_frame_action: String,
#[arg(
long,
default_value = "frame_diff",
help = "off | frame_diff — inter-frame motion proxy"
)]
motion_mode: String,
#[arg(long, default_value_t = 8, help = "Frames per GPU batch (throughput)")]
batch: u32,
}
impl ExtractOpts {
fn to_config(&self) -> Result<ExtractConfig, PipelineError> {
Ok(ExtractConfig {
chunk_size: self.chunk_size,
scale_width: self.scale_width,
workers: self.workers,
output_pattern: self.output_pattern.clone(),
output_format: self.output_format.clone(),
jpeg_quality: self.jpeg_quality,
reuse_metadata: !self.reuse_metadata_inv,
min_sharpness: self.min_sharpness,
min_exposure: self.min_exposure,
max_exposure: self.max_exposure,
min_contrast: self.min_contrast,
min_motion: self.min_motion,
max_motion: self.max_motion,
selection: Selection::parse(&self.selection)
.map_err(|e| PipelineError(e.to_string()))?,
target_count: self.target_count,
min_gap: self.min_gap,
bad_frame_action: BadFrameAction::parse(&self.bad_frame_action)
.map_err(|e| PipelineError(e.to_string()))?,
motion_mode: MotionMode::parse(&self.motion_mode)
.map_err(|e| PipelineError(e.to_string()))?,
batch: self.batch,
})
}
}
// Shared extract flags for the `all` subcommand, with `--extract-*` prefixes
/// (mirrors the Python `all` CLI the GUI drives).
#[derive(Debug, Args, Clone)]
struct ExtractOptsAll {
#[arg(long = "extract-chunk-size", default_value_t = 30)]
chunk_size: u32,
#[arg(long = "extract-scale-width", default_value_t = 1920)]
scale_width: u32,
#[arg(long = "extract-workers", default_value_t = 4)]
workers: u32,
#[arg(
long = "extract-output-pattern",
default_value = "output_frame_%05d.png"
)]
output_pattern: String,
#[arg(long = "extract-output-format", default_value = "png")]
output_format: String,
#[arg(long = "extract-jpeg-quality", default_value_t = 95)]
jpeg_quality: u32,
#[arg(long = "extract-regenerate-metadata")]
reuse_metadata_inv: bool,
#[arg(long = "extract-min-sharpness", default_value_t = 0.0)]
min_sharpness: f64,
#[arg(long = "extract-min-exposure", default_value_t = 0.0)]
min_exposure: f64,
#[arg(long = "extract-max-exposure", default_value_t = 255.0)]
max_exposure: f64,
#[arg(long = "extract-min-contrast", default_value_t = 0.0)]
min_contrast: f64,
#[arg(long = "extract-min-motion", default_value_t = 0.0)]
min_motion: f64,
#[arg(long = "extract-max-motion", default_value_t = 1.0e12)]
max_motion: f64,
#[arg(long = "extract-selection", default_value = "chunk")]
selection: String,
#[arg(long = "extract-target-count", default_value_t = 0)]
target_count: u32,
#[arg(long = "extract-min-gap", default_value_t = 0)]
min_gap: u32,
#[arg(long = "extract-bad-frame-action", default_value = "keep")]
bad_frame_action: String,
#[arg(long = "extract-motion-mode", default_value = "frame_diff")]
motion_mode: String,
#[arg(long = "extract-batch", default_value_t = 8)]
batch: u32,
}
impl ExtractOptsAll {
fn to_config(&self) -> Result<ExtractConfig, PipelineError> {
Ok(ExtractConfig {
chunk_size: self.chunk_size,
scale_width: self.scale_width,
workers: self.workers,
output_pattern: self.output_pattern.clone(),
output_format: self.output_format.clone(),
jpeg_quality: self.jpeg_quality,
reuse_metadata: !self.reuse_metadata_inv,
min_sharpness: self.min_sharpness,
min_exposure: self.min_exposure,
max_exposure: self.max_exposure,
min_contrast: self.min_contrast,
min_motion: self.min_motion,
max_motion: self.max_motion,
selection: Selection::parse(&self.selection)
.map_err(|e| PipelineError(e.to_string()))?,
target_count: self.target_count,
min_gap: self.min_gap,
bad_frame_action: BadFrameAction::parse(&self.bad_frame_action)
.map_err(|e| PipelineError(e.to_string()))?,
motion_mode: MotionMode::parse(&self.motion_mode)
.map_err(|e| PipelineError(e.to_string()))?,
batch: self.batch,
})
}
}
#[derive(Debug, Args)]
struct AllArgs {
#[arg(long)]
project_dir: PathBuf,
#[arg(long, num_args = 1..)]
video: Vec<PathBuf>,
#[arg(long, default_value = "extract")]
start_stage: String,
#[arg(long)]
skip_masks: bool,
#[arg(long, hide = true)]
gui_log_output: bool,
#[command(flatten)]
extract: ExtractOptsAll,
#[command(flatten)]
mask: MasksOptsAll,
#[arg(long = "cubemap-face-size", default_value_t = 1600)]
cubemap_face_size: u32,
#[arg(long = "cubemap-skip-faces", default_value = "")]
cubemap_skip_faces: String,
#[arg(long = "cubemap-jpeg-quality", default_value_t = 95)]
cubemap_jpeg_quality: u32,
#[arg(long = "cubemap-batch", default_value_t = 4)]
cubemap_batch: u32,
#[arg(long = "rig-ref-face", default_value = "F")]
rig_ref_face: String,
#[arg(long = "rig-intrinsics-convention", default_value = "linspace")]
rig_intrinsics_convention: String,
#[arg(long = "rig-link-mode", default_value = "symlink")]
rig_link_mode: String,
}
#[derive(Debug, Args)]
struct ExtractArgs {
#[arg(long, num_args = 1..)]
video: Vec<PathBuf>,
#[arg(long)]
output_dir: PathBuf,
#[arg(long, hide = true)]
gui_log_output: bool,
#[command(flatten)]
extract: ExtractOpts,
}
// Shared masks-stage flags for the `masks` subcommand (unprefixed).
#[derive(Debug, Args, Clone)]
struct MasksOpts {
#[arg(long, default_value = "person,car,bus,truck,bicycle,motorcycle")]
classes: String,
#[arg(
long,
default_value = "yolo11x-seg.onnx",
help = "ONNX segmentation model path"
)]
model: String,
#[arg(
long,
default_value = "cpu",
help = "ort execution device (cpu|cuda|rocm)"
)]
device: String,
#[arg(long, default_value_t = 0.35)]
conf: f32,
#[arg(long, default_value_t = 1280)]
imgsz: u32,
#[arg(long, default_value_t = 6)]
dilate: u32,
#[arg(long, default_value = "cubemap", help = "cubemap | equirect")]
projection: String,
#[arg(long, default_value_t = 1024)]
face_size: u32,
#[arg(long, default_value = "")]
skip_faces: String,
#[arg(long, default_value_t = 8)]
mask_batch: u32,
#[arg(long, default_value_t = 8)]
tile_batch: u32,
#[arg(long = "mask-empty-if-none")]
empty_if_none: bool,
#[arg(long, default_value_t = 0.45, help = "NMS IoU threshold")]
iou: f32,
#[arg(
long,
help = "Directory for the MIGraphX compiled-model cache (GPU path). First \
run compiles+saves; later runs load. Pair with --device rocm. \
Defaults to ~/.cache/pipeline_360/migraphx."
)]
ep_cache: Option<PathBuf>,
}
impl MasksOpts {
fn to_config(&self) -> Result<MasksConfig, PipelineError> {
let skip = cubemap::parse_skip_faces(&self.skip_faces)
.map_err(|e| PipelineError(e.to_string()))?;
Ok(MasksConfig {
classes: masks::resolve_classes(&self.classes)
.map_err(|e| PipelineError(e.to_string()))?,
model: PathBuf::from(&self.model),
device: self.device.clone(),
conf: self.conf,
imgsz: self.imgsz,
dilate: self.dilate,
projection: Projection::parse(&self.projection)
.map_err(|e| PipelineError(e.to_string()))?,
face_size: self.face_size,
skip_faces: skip,
batch: self.mask_batch,
tile_batch: self.tile_batch,
empty_if_none: self.empty_if_none,
iou: self.iou,
ep_cache: self.ep_cache.clone(),
})
}
}
// Shared masks-stage flags for `all` (with `--mask-*` prefixes, matching Python).
#[derive(Debug, Args, Clone)]
struct MasksOptsAll {
#[arg(
long = "mask-classes",
default_value = "person,car,bus,truck,bicycle,motorcycle"
)]
classes: String,
#[arg(long = "mask-model", default_value = "yolo11x-seg.onnx")]
model: String,
#[arg(long = "mask-device", default_value = "cpu")]
device: String,
#[arg(long = "mask-conf", default_value_t = 0.35)]
conf: f32,
#[arg(long = "mask-imgsz", default_value_t = 1280)]
imgsz: u32,
#[arg(long = "mask-dilate", default_value_t = 6)]
dilate: u32,
#[arg(long = "mask-projection", default_value = "cubemap")]
projection: String,
#[arg(long = "mask-face-size", default_value_t = 1024)]
face_size: u32,
#[arg(long = "mask-skip-faces", default_value = "")]
skip_faces: String,
#[arg(long = "mask-batch", default_value_t = 8)]
mask_batch: u32,
#[arg(long = "mask-tile-batch", default_value_t = 8)]
mask_tile_batch: u32,
#[arg(long = "mask-empty-if-none")]
empty_if_none: bool,
#[arg(long = "mask-iou", default_value_t = 0.45)]
iou: f32,
#[arg(
long = "mask-ep-cache",
help = "Directory for the MIGraphX compiled-model cache (GPU path). First \
run compiles+saves; later runs load. Pair with --mask-device rocm. \
Defaults to ~/.cache/pipeline_360/migraphx."
)]
ep_cache: Option<PathBuf>,
}
impl MasksOptsAll {
fn to_config(&self) -> Result<MasksConfig, PipelineError> {
let skip = cubemap::parse_skip_faces(&self.skip_faces)
.map_err(|e| PipelineError(e.to_string()))?;
Ok(MasksConfig {
classes: masks::resolve_classes(&self.classes)
.map_err(|e| PipelineError(e.to_string()))?,
model: PathBuf::from(&self.model),
device: self.device.clone(),
conf: self.conf,
imgsz: self.imgsz,
dilate: self.dilate,
projection: Projection::parse(&self.projection)
.map_err(|e| PipelineError(e.to_string()))?,
face_size: self.face_size,
skip_faces: skip,
batch: self.mask_batch,
tile_batch: self.mask_tile_batch,
empty_if_none: self.empty_if_none,
iou: self.iou,
ep_cache: self.ep_cache.clone(),
})
}
}
#[derive(Debug, Args)]
struct MasksArgs {
#[arg(long)]
input_dir: PathBuf,
#[arg(long)]
output_dir: PathBuf,
#[arg(long, hide = true)]
gui_log_output: bool,
#[command(flatten)]
opts: MasksOpts,
}
#[derive(Debug, Args)]
struct CubemapArgs {
#[arg(long)]
frames_dir: PathBuf,
#[arg(long)]
masks_dir: Option<PathBuf>,
#[arg(long)]
output_frames_dir: PathBuf,
#[arg(long)]
output_masks_dir: Option<PathBuf>,
#[arg(long, default_value_t = 1600)]
face_size: u32,
#[arg(long, default_value = "")]
skip_faces: String,
#[arg(long, default_value_t = 95)]
jpeg_quality: u32,
#[arg(long)]
resume: bool,
#[arg(long, default_value_t = 4, help = "Frames per GPU batch (throughput)")]
cubemap_batch: u32,
#[arg(long, hide = true)]
gui_log_output: bool,
}
#[derive(Debug, Args)]
struct RigArgs {
#[arg(long)]
images_dir: Option<PathBuf>,
#[arg(long)]
masks_dir: Option<PathBuf>,
#[arg(long, default_value = "rig_config.json")]
output: PathBuf,
#[arg(long)]
face_size: Option<u32>,
#[arg(long, default_value = "auto")]
layout: String,
#[arg(long)]
link_tree: Option<PathBuf>,
#[arg(long)]
masks_link_tree: Option<PathBuf>,
#[arg(long, default_value = "symlink")]
link_mode: String,
#[arg(long, default_value = "F")]
ref_face: String,
#[arg(long, default_value = "linspace")]
intrinsics_convention: String,
#[arg(long)]
self_test: bool,
#[arg(long, hide = true)]
gui_log_output: bool,
}
fn run() -> Result<(), PipelineError> {
let cli = Cli::parse();
let sub_gui = match &cli.command {
Some(Command::All(a)) => a.gui_log_output,
Some(Command::Extract(a)) => a.gui_log_output,
Some(Command::Masks(a)) => a.gui_log_output,
Some(Command::Cubemap(a)) => a.gui_log_output,
Some(Command::Rig(a)) => a.gui_log_output,
None => false,
};
let gui_log = cli.gui_log_output || sub_gui;
let logger: Logger = if gui_log {
Some(&|m| println!("{m}"))
} else {
None
};
let cancel = CancelFlag::new();
// No subcommand (or explicit --gui): launch the native GUI if available.
if cli.command.is_none() || cli.gui {
#[cfg(feature = "gui")]
{
pipeline_360::gui::launch().map_err(|e| PipelineError(format!("GUI: {e}")))?;
return Ok(());
}
#[cfg(not(feature = "gui"))]
{
use clap::CommandFactory;
Cli::command()
.print_help()
.map_err(|e| PipelineError(format!("print_help: {e}")))?;
println!("\n\n(no subcommand given; rebuild with `--features gui` to launch the GUI)");
return Ok(());
}
}
match cli.command {
Some(Command::Rig(RigArgs {
self_test: true, ..
})) => {
let ok = pipeline_360::selftest::run_selftest().map_err(PipelineError)?;
logging::emit(&format!("self-test: {ok} check groups passed."), logger);
}
Some(Command::Rig(a)) => {
let images_dir = a.images_dir.as_deref().ok_or_else(|| {
PipelineError("--images-dir is required (or use --self-test)".into())
})?;
if a.masks_link_tree.is_some() && a.masks_dir.is_none() {
return Err(PipelineError(
"--masks-link-tree requires --masks-dir".into(),
));
}
run_rig(images_dir, &a, logger)?;
}
Some(Command::All(a)) => run_all(&a, logger, &cancel)?,
Some(Command::Extract(a)) => {
let cfg = a.extract.to_config()?;
let gpu = pollster::block_on(GpuContext::new("extract"))
.map_err(|e| PipelineError(format!("GPU init: {e}")))?;
extract::run_extract_stage(&a.video, &a.output_dir, &cfg, &gpu, &cancel, logger)
.map_err(|e| PipelineError(e.to_string()))?;
}
Some(Command::Masks(a)) => {
let cfg = a.opts.to_config()?;
let gpu = pollster::block_on(GpuContext::new("masks"))
.map_err(|e| PipelineError(format!("GPU init: {e}")))?;
masks::run_masks_stage(&a.input_dir, &a.output_dir, &cfg, &gpu, &cancel, logger)
.map_err(|e| PipelineError(e.to_string()))?;
}
Some(Command::Cubemap(a)) => run_cubemap(a, &cancel, logger)?,
None => unreachable!(),
}
Ok(())
}
fn run_cubemap(a: CubemapArgs, cancel: &CancelFlag, logger: Logger) -> Result<(), PipelineError> {
if a.masks_dir.is_some() && a.output_masks_dir.is_none() {
return Err(PipelineError(
"--masks-dir given but --output-masks-dir missing".into(),
));
}
let gpu = pollster::block_on(GpuContext::new("cubemap"))
.map_err(|e| PipelineError(format!("GPU init: {e}")))?;
let skip =
cubemap::parse_skip_faces(&a.skip_faces).map_err(|e| PipelineError(e.to_string()))?;
let (fout, mout) = cubemap::run_cubemap_stage(
&a.frames_dir,
&a.output_frames_dir,
a.masks_dir.as_deref(),
a.output_masks_dir.as_deref(),
a.face_size,
&skip,
a.jpeg_quality,
a.resume,
a.cubemap_batch,
&gpu,
cancel,
logger,
)
.map_err(|e| PipelineError(e.to_string()))?;
logging::emit(&format!("Cubemap output: {}", fout.display()), logger);
if let Some(m) = mout {
logging::emit(&format!("Cubemap masks: {}", m.display()), logger);
}
Ok(())
}
fn run_rig(images_dir: &std::path::Path, a: &RigArgs, logger: Logger) -> Result<(), PipelineError> {
let layout = match a.layout.as_str() {
"flat" => Layout::Flat,
"nested" => Layout::Nested,
_ => Layout::Auto,
};
let mode = LinkMode::parse(&a.link_mode)?;
let out = rig::run_rig_stage(
images_dir,
&a.output,
a.masks_dir.as_deref(),
a.face_size,
layout,
&a.ref_face,
&a.intrinsics_convention,
a.link_tree.as_deref(),
a.masks_link_tree.as_deref(),
mode,
)?;
for msg in &out.messages {
logging::emit(msg, logger);
}
Ok(())
}
fn run_all(a: &AllArgs, logger: Logger, cancel: &CancelFlag) -> Result<(), PipelineError> {
let project = ProjectLayout::new(&a.project_dir);
project
.create_all()
.map_err(|e| PipelineError(format!("mkdir: {e}")))?;
let stage_idx = |s: &str| match s {
"extract" => 0,
"cubemap" => 1,
"masks" => 2,
"rig" => 3,
_ => 0,
};
let from = stage_idx(&a.start_stage);
// NOTE: order is extract -> cubemap -> masks -> rig. Masks run on the
// cube-face images the cubemap stage produces, so cubemap must come first
// (masking equirect frames is useless, and round-tripping e2c->YOLO->c2e
// then e2c again was wasteful). `02_masks` is no longer written by `all`.
let stages = ["extract", "cubemap", "masks", "rig"];
if a.video.is_empty() && from == 0 {
return Err(PipelineError(
"--video is required when starting from the extract stage.".into(),
));
}
// One GPU device shared by every GPU stage in the run.
let gpu = pollster::block_on(GpuContext::new("pipeline"))
.map_err(|e| PipelineError(format!("GPU init: {e}")))?;
for (i, stage) in stages.iter().enumerate() {
if i < from {
continue;
}
cancel
.check()
.map_err(|_| PipelineError("Cancelled.".into()))?;
logging::stage_header(i + 1, 4, stage_title(stage), logger);
match *stage {
"extract" => {
let cfg = a.extract.to_config()?;
extract::run_extract_stage(&a.video, &project.frames, &cfg, &gpu, cancel, logger)
.map_err(|e| PipelineError(e.to_string()))?;
}
"masks" => {
if a.skip_masks {
logging::emit("Skipping masks stage (--skip-masks).", logger);
} else {
let cfg = a.mask.to_config()?;
// Masks run on the cube-face images the cubemap stage
// produced, and write cube-face masks next to them.
masks::run_masks_stage(
&project.cubemap_images,
&project.cubemap_masks,
&cfg,
&gpu,
cancel,
logger,
)
.map_err(|e| PipelineError(e.to_string()))?;
}
}
"cubemap" => {
let skip = cubemap::parse_skip_faces(&a.cubemap_skip_faces)
.map_err(|e| PipelineError(e.to_string()))?;
// In the `all` flow the cubemap stage only projects images; the
// masks stage runs afterwards on those cube faces (no equirect
// mask round-trip, no double e2c projection).
cubemap::run_cubemap_stage(
&project.frames,
&project.cubemap_images,
None,
None,
a.cubemap_face_size,
&skip,
a.cubemap_jpeg_quality,
false,
a.cubemap_batch,
&gpu,
cancel,
logger,
)
.map_err(|e| PipelineError(e.to_string()))?;
}
"rig" => {
let mode = LinkMode::parse(&a.rig_link_mode)?;
let out = rig::run_rig_stage(
&project.cubemap_images,
&project.rig_config,
if a.skip_masks {
None
} else {
Some(project.cubemap_masks.as_path())
},
None,
Layout::Auto,
&a.rig_ref_face,
&a.rig_intrinsics_convention,
Some(&project.colmap_images),
if a.skip_masks {
None
} else {
Some(&project.colmap_masks)
},
mode,
)
.map_err(|e| PipelineError(e.to_string()))?;
for msg in &out.messages {
logging::emit(msg, logger);
}
}
_ => {}
}
}
logging::emit("Pipeline complete.", logger);
Ok(())
}
fn stage_title(stage: &str) -> &'static str {
match stage {
"extract" => "Extract Sharpest Frames",
"masks" => "Generate SfM Masks",
"cubemap" => "Project to Cubemap",
"rig" => "Generate Rig Config",
_ => "Unknown",
}
}
fn main() -> ExitCode {
match run() {
Ok(()) => ExitCode::SUCCESS,
Err(e) => {
eprintln!("Error: {e}");
ExitCode::FAILURE
}
}
}
+172
View File
@@ -0,0 +1,172 @@
//! Rig geometry math — a direct port of the pure-logic core of the Python
//! `pipeline_360.py` stage-4 geometry (`FACE_AXES`, `compute_intrinsics`,
//! `cam_from_rig_rotation`, …).
//!
//! This module is deliberately **pure `std`** (no external crates) so it can be
//! compiled and tested offline with plain `rustc`, and so the geometry core can
//! be reused verbatim across every stage's GPU kernels.
//!
//! Conventions (verified against the Python stage-4 self-test and
//! py360convert source):
//! * +X = right, +Y = up, +Z = front in an endpoint-inclusive linspace grid.
//! * fx = fy = (N-1)/2 for `linspace` (exact for py360convert), N/2 for `half`.
//! * COLMAP `cam_from_rig_rotation` is W-FIRST ((w,x,y,z) quaternion).
/// Face letters in canonical order.
pub const FACES: [&str; 6] = ["F", "R", "B", "L", "U", "D"];
/// Human-readable face names (mirrors Python `FACE_NAMES`).
pub const FACE_NAMES: [(&str, &str); 6] = [
("F", "front"),
("R", "right"),
("B", "back"),
("L", "left"),
("U", "up"),
("D", "down"),
];
/// A 3-vector.
pub type Vec3 = [f64; 3];
/// A row-major 3x3 matrix.
pub type Mat3 = [[f64; 3]; 3];
/// Face axes: `(right, down, forward)` — image right, image down, view direction.
///
/// Ported from Python `FACE_AXES`. `forward` is the view direction, `right` the
/// +u image axis and `down` the +v image axis in cube-face space.
pub fn face_rotation_matrix(face: &str) -> Mat3 {
match face {
"F" => [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
"R" => [[0.0, 0.0, -1.0], [0.0, 1.0, 0.0], [1.0, 0.0, 0.0]],
"B" => [[-1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, -1.0]],
"L" => [[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]],
"U" => [[1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, -1.0, 0.0]],
"D" => [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]],
_ => panic!("unknown face: {face}"),
}
}
/// Matrix multiply: `a @ b`.
pub fn mat_mul(a: &Mat3, b: &Mat3) -> Mat3 {
let mut out = Mat3::default();
for i in 0..3 {
for j in 0..3 {
out[i][j] = (0..3).map(|k| a[i][k] * b[k][j]).sum();
}
}
out
}
/// Transpose.
pub fn mat_transpose(m: &Mat3) -> Mat3 {
let mut out = Mat3::default();
for i in 0..3 {
for j in 0..3 {
out[j][i] = m[i][j];
}
}
out
}
/// Determinant.
pub fn mat_det(m: &Mat3) -> f64 {
m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
- m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
+ m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0])
}
/// Assert `m` is a proper rotation (det == +1 and R @ R^T == I).
pub fn assert_valid_rotation(m: &Mat3, label: &str, tol: f64) -> Result<(), String> {
let det = mat_det(m);
if (det - 1.0).abs() > tol {
return Err(format!("{label}: det={det}, expected +1 (not a proper rotation)"));
}
let ident = mat_mul(m, &mat_transpose(m));
for i in 0..3 {
for j in 0..3 {
let expected = if i == j { 1.0 } else { 0.0 };
if (ident[i][j] - expected).abs() > tol {
return Err(format!("{label}: R @ R.T is not identity ({ident:?})"));
}
}
}
Ok(())
}
/// Apply matrix to vector (row-major): `m @ v`.
pub fn apply(m: &Mat3, v: &Vec3) -> Vec3 {
let mut out = [0.0; 3];
for i in 0..3 {
out[i] = m[i][0] * v[0] + m[i][1] * v[1] + m[i][2] * v[2];
}
out
}
/// Shepperd's method → COLMAP's `(w, x, y, z)` quaternion, sign-canonicalised so `w >= 0`.
pub fn quaternion_from_matrix(m: &Mat3) -> (f64, f64, f64, f64) {
let m00 = m[0][0];
let m01 = m[0][1];
let m02 = m[0][2];
let m10 = m[1][0];
let m11 = m[1][1];
let m12 = m[1][2];
let m20 = m[2][0];
let m21 = m[2][1];
let m22 = m[2][2];
let trace = m00 + m11 + m22;
let (w, x, y, z) = if trace > 0.0 {
let s = (trace + 1.0).sqrt() * 2.0;
(0.25 * s, (m21 - m12) / s, (m02 - m20) / s, (m10 - m01) / s)
} else if m00 > m11 && m00 > m22 {
let s = (1.0 + m00 - m11 - m22).sqrt() * 2.0;
((m21 - m12) / s, 0.25 * s, (m01 + m10) / s, (m02 + m20) / s)
} else if m11 > m22 {
let s = (1.0 + m11 - m00 - m22).sqrt() * 2.0;
((m02 - m20) / s, (m01 + m10) / s, 0.25 * s, (m12 + m21) / s)
} else {
let s = (1.0 + m22 - m00 - m11).sqrt() * 2.0;
((m10 - m01) / s, (m02 + m20) / s, (m12 + m21) / s, 0.25 * s)
};
let norm = (w * w + x * x + y * y + z * z).sqrt();
let mut q = (w / norm, x / norm, y / norm, z / norm);
if q.0 < 0.0 {
q = (-q.0, -q.1, -q.2, -q.3);
}
q
}
/// The rotation of `face` relative to `ref_face`, as a W-first COLMAP quaternion.
pub fn cam_from_rig_rotation(face: &str, ref_face: &str) -> (f64, f64, f64, f64) {
let r_face = face_rotation_matrix(face);
let r_ref = face_rotation_matrix(ref_face);
assert_valid_rotation(&r_face, &format!("FACE_AXES[{face}]"), 1e-9)
.expect("invalid FACE_AXES");
assert_valid_rotation(&r_ref, &format!("FACE_AXES[{ref_face}]"), 1e-9)
.expect("invalid FACE_AXES");
let r = mat_mul(&r_face, &mat_transpose(&r_ref));
quaternion_from_matrix(&r)
}
/// Rotation angle in degrees implied by a quaternion (to validate ±90/±180 steps).
pub fn rotation_angle_deg(q: (f64, f64, f64, f64)) -> f64 {
let w = q.0.abs().min(1.0).max(-1.0);
2.0 * w.acos() * 180.0 / std::f64::consts::PI
}
/// `(fx, fy, cx, cy)`. `linspace`: f=(N-1)/2 (exact for py360convert); `half`: f=N/2.
pub fn compute_intrinsics(face_size: u32, convention: &str) -> (f64, f64, f64, f64) {
if face_size < 2 {
panic!("face_size must be >= 2, got {face_size}");
}
let n = f64::from(face_size);
let f = match convention {
"linspace" => (n - 1.0) / 2.0,
"half" => n / 2.0,
other => panic!("unknown intrinsics convention: {other}"),
};
let c = n / 2.0;
(f, f, c, c)
}
+74
View File
@@ -0,0 +1,74 @@
//! Project directory layout + file gathering (pure `std`).
use std::path::{Path, PathBuf};
/// Stage order — cubemap runs *before* masks here (masks run on the cube-face
/// images the cubemap stage produces, so there is no equirect round-trip).
pub const STAGE_ORDER: [&str; 4] = ["extract", "cubemap", "masks", "rig"];
/// Image extensions considered by `gather_images` (mirrors Python).
pub const IMAGE_EXTENSIONS: [&str; 7] = [".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"];
/// Resolve the standard sub-directory names for a project root.
pub struct ProjectLayout {
pub root: PathBuf,
pub frames: PathBuf, // 01_frames
pub masks: PathBuf, // 02_masks
pub cubemap_images: PathBuf, // 03_cubemap/images
pub cubemap_masks: PathBuf, // 03_cubemap/masks
pub colmap_images: PathBuf, // 04_colmap/images
pub colmap_masks: PathBuf, // 04_colmap/masks
pub rig_config: PathBuf, // 04_colmap/rig_config.json
}
impl ProjectLayout {
pub fn new(root: impl Into<PathBuf>) -> Self {
let r = root.into();
ProjectLayout {
cubemap_images: r.join("03_cubemap").join("images"),
cubemap_masks: r.join("03_cubemap").join("masks"),
colmap_images: r.join("04_colmap").join("images"),
colmap_masks: r.join("04_colmap").join("masks"),
masks: r.join("02_masks"),
frames: r.join("01_frames"),
rig_config: r.join("04_colmap").join("rig_config.json"),
root: r,
}
}
pub fn create_all(&self) -> std::io::Result<()> {
std::fs::create_dir_all(&self.root)?;
std::fs::create_dir_all(&self.frames)?;
std::fs::create_dir_all(&self.masks)?;
std::fs::create_dir_all(&self.cubemap_images)?;
std::fs::create_dir_all(&self.cubemap_masks)?;
Ok(())
}
}
/// Sort files by name; returns an error if `dir` has no image files.
pub fn gather_images(input_dir: &Path) -> Result<Vec<PathBuf>, String> {
let ok = std::fs::read_dir(input_dir)
.map_err(|e| format!("read dir {}: {e}", input_dir.display()))?;
let mut frames: Vec<PathBuf> = ok
.flatten()
.map(|e| e.path())
.filter(|p| is_image(p))
.collect();
frames.sort();
if frames.is_empty() {
return Err(format!("No image files found in {}", input_dir.display()));
}
Ok(frames)
}
/// Whether a path has one of the known image extensions.
pub fn is_image(p: &Path) -> bool {
p.extension()
.and_then(|e| e.to_str())
.map(|e| {
let e = e.to_lowercase();
IMAGE_EXTENSIONS.iter().any(|x| *x == format!(".{e}"))
})
.unwrap_or(false)
}
+326
View File
@@ -0,0 +1,326 @@
//! Port of `pipeline_360.py rig --self-test`: validates stage-4 geometry &
//! schema without any video or GPU. Pure `std` so it can run offline.
use std::path::Path;
use crate::json;
use crate::math;
use crate::stages::rig::{self, Layout, LinkMode};
/// Write a tiny valid (stored/deflate) single-colour PNG so rig's pure-std
/// dimension reader and copy-linker can operate on it without a codec.
fn write_dummy_png(path: &Path, width: u32, height: u32) -> std::io::Result<()> {
// Build an uncompressed-all-zero RGB image: each scanline prefixed with a
// filter byte 0, then width RGB triples.
let mut raw: Vec<u8> = Vec::with_capacity((height as usize) * (1 + 3 * width as usize));
for _ in 0..height {
raw.push(0);
raw.extend(std::iter::repeat(0u8).take((3 * width) as usize));
}
// zlib stream: a single "stored" deflate block (BTYPE=00) with the full
// uncompressed length, no compression. Valid zlib with header 0x78 0x01.
let len = raw.len() as u32;
let mut deflate = Vec::new();
deflate.push(0x01); // BFINAL=1, BTYPE=00
deflate.extend_from_slice(&len.to_le_bytes());
deflate.extend_from_slice(&(!len).to_le_bytes());
deflate.extend_from_slice(&raw);
let mut idat = vec![0x78, 0x01];
// zlib adler32 of `raw`.
let adler = adler32(&raw);
idat.extend_from_slice(&deflate);
idat.extend_from_slice(&adler.to_be_bytes());
let ihdr_data = make_ihdr(width, height);
let mut out = Vec::new();
out.extend_from_slice(b"\x89PNG\r\n\x1a\n");
png_chunk(&mut out, b"IHDR", &ihdr_data);
png_chunk(&mut out, b"IDAT", &idat);
png_chunk(&mut out, b"IEND", &[]);
std::fs::create_dir_all(path.parent().unwrap_or(Path::new(".")))?;
std::fs::write(path, out)
}
fn make_ihdr(width: u32, height: u32) -> Vec<u8> {
let mut d = Vec::new();
d.extend_from_slice(&width.to_be_bytes());
d.extend_from_slice(&height.to_be_bytes());
d.extend_from_slice(&[8, 2, 0, 0, 0]); // bit depth 8, colour type 2 (RGB)
d
}
fn png_chunk(out: &mut Vec<u8>, tag: &[u8; 4], data: &[u8]) {
out.extend_from_slice(&(data.len() as u32).to_be_bytes());
out.extend_from_slice(tag);
out.extend_from_slice(data);
// CRC32 of (tag + data).
let mut crc_input = Vec::new();
crc_input.extend_from_slice(tag);
crc_input.extend_from_slice(data);
out.extend_from_slice(&crc32(&crc_input).to_be_bytes());
}
fn crc32(data: &[u8]) -> u32 {
let mut table = [0u32; 256];
for (n, item) in table.iter_mut().enumerate() {
let mut c = n as u32;
for _ in 0..8 {
c = if c & 1 != 0 {
0xEDB8_8320 ^ (c >> 1)
} else {
c >> 1
};
}
*item = c;
}
let mut c = 0xFFFF_FFFFu32;
for &b in data {
c = table[((c ^ b as u32) & 0xFF) as usize] ^ (c >> 8);
}
c ^ 0xFFFF_FFFF
}
fn adler32(data: &[u8]) -> u32 {
let mut a = 1u32;
let mut b = 0u32;
for &byte in data {
a = (a + byte as u32) % 65521;
b = (b + a) % 65521;
}
(b << 16) | a
}
fn approx_eq_vec(a: &[f64], b: &[f64], tol: f64) -> bool {
a.len() == b.len() && a.iter().zip(b).all(|(x, y)| (x - y).abs() <= tol)
}
/// Run the six check groups. Returns `(ok_count, first_error)`.
pub fn run_selftest() -> Result<usize, String> {
let n = 64u32;
let mut ok = 0usize;
// Group 1: every FACE_AXES rotation is valid.
for face in math::FACES {
let m = math::face_rotation_matrix(face);
math::assert_valid_rotation(&m, &format!("FACE_AXES[{face}]"), 1e-9)
.map_err(|e| format!("group 1: {e}"))?;
}
ok += 1;
// Group 2: R<-F quaternion and 90° angle.
let q_r = math::cam_from_rig_rotation("R", "F");
let sqrt2_2 = std::f64::consts::SQRT_2 / 2.0;
let expected_r = [sqrt2_2, 0.0, -sqrt2_2, 0.0];
if !approx_eq_vec(&[q_r.0, q_r.1, q_r.2, q_r.3], &expected_r, 1e-12) {
return Err(format!("group 2: R quaternion wrong: {q_r:?}"));
}
if (math::rotation_angle_deg(q_r) - 90.0).abs() >= 1e-9 {
return Err("group 2: rotation angle not 90°".into());
}
ok += 1;
// Group 3: forward directions from each face's transpose mapping +Z.
let r_mat = math::face_rotation_matrix("R");
let fwd_f = math::apply(&math::mat_transpose(&r_mat), &[0.0, 0.0, 1.0]);
if !approx_eq_vec(&fwd_f, &[1.0, 0.0, 0.0], 1e-12) {
return Err(format!("group 3: R forward wrong: {fwd_f:?}"));
}
for (face, expected) in [
("B", [0.0, 0.0, -1.0]),
("L", [-1.0, 0.0, 0.0]),
("U", [0.0, -1.0, 0.0]),
("D", [0.0, 1.0, 0.0]),
] {
let got = math::apply(
&math::mat_transpose(&math::face_rotation_matrix(face)),
&[0.0, 0.0, 1.0],
);
if !approx_eq_vec(&got, &expected, 1e-12) {
return Err(format!("group 3: face {face} forward wrong: {got:?}"));
}
}
ok += 1;
// Group 4: intrinsics.
let (fx, fy, cx, cy) = math::compute_intrinsics(n, "linspace");
if !approx_eq_vec(
&[fx, fy, cx, cy],
&[
(n - 1) as f64 / 2.0,
(n - 1) as f64 / 2.0,
n as f64 / 2.0,
n as f64 / 2.0,
],
0.0,
) {
return Err(format!("group 4: intrinsics wrong: {fx},{fy},{cx},{cy}"));
}
ok += 1;
// Group 5: full rig run on a built-up dataset.
let tmp = std::env::temp_dir().join(format!("rig_selftest_{}", std::process::id()));
let root = tmp.as_path();
let _ = std::fs::remove_dir_all(root);
let flat = root.join("images");
let masks_flat = root.join("masks");
for i in 1..=2u32 {
for face in math::FACES {
write_dummy_png(&flat.join(format!("frame_{i:05}_{face}.png")), n, n)
.map_err(|e| format!("group 5: write image: {e}"))?;
}
}
for face in ["F", "R", "B", "L"] {
write_dummy_png(&flat.join(format!("frame_00003_{face}.png")), n, n)
.map_err(|e| format!("group 5: write partial image: {e}"))?;
}
for face in math::FACES {
write_dummy_png(
&masks_flat.join(format!("frame_00001_{face}.png.png")),
n,
n,
)
.map_err(|e| format!("group 5: write mask: {e}"))?;
}
for face in ["F", "R"] {
write_dummy_png(
&masks_flat.join(format!("frame_00003_{face}.png.png")),
n,
n,
)
.map_err(|e| format!("group 5: write partial mask: {e}"))?;
}
let out = root.join("rig.json");
let link_tree = root.join("nested");
let masks_link_tree = root.join("masks_nested");
let _res = rig::run_rig_stage(
&flat,
&out,
Some(&masks_flat),
None,
Layout::Auto,
"F",
"linspace",
Some(&link_tree),
Some(&masks_link_tree),
LinkMode::Copy,
)
.map_err(|e| format!("group 5: run_rig_stage: {e}"))?;
let text = std::fs::read_to_string(&out).map_err(|e| format!("group 5: read rig: {e}"))?;
let config = json::parse(&text).map_err(|e| format!("group 5: parse rig json: {e}"))?;
let cams = config
.as_array()
.and_then(|a| a.first())
.and_then(|v| v.find("cameras"))
.and_then(|v| v.as_array())
.ok_or("group 5: no cameras array")?;
if cams.len() != 6 {
return Err(format!("group 5: expected 6 cameras, got {}", cams.len()));
}
let refs: Vec<&json::Value> = cams
.iter()
.filter(|c| matches!(c.find("ref_sensor"), Some(crate::json::Value::Bool(true))))
.collect();
if refs.len() != 1 {
return Err(format!(
"group 5: expected exactly 1 ref_sensor, got {}",
refs.len()
));
}
if refs[0].find("image_prefix").and_then(|v| v.as_str()) != Some("F/") {
return Err("group 5: ref prefix not F/".into());
}
let mut by_prefix: std::collections::HashMap<&str, &json::Value> =
std::collections::HashMap::new();
for c in cams {
if let Some(p) = c.find("image_prefix").and_then(|v| v.as_str()) {
by_prefix.insert(p, c);
}
}
let got_r: Vec<f64> = by_prefix
.get("R/")
.and_then(|c| c.find("cam_from_rig_rotation"))
.and_then(|v| v.as_array())
.map(|a| a.iter().filter_map(|x| x.as_f64()).collect())
.ok_or("group 5: no R rotation")?;
if !approx_eq_vec(&got_r, &expected_r, 1e-12) {
return Err(format!("group 5: R rotation wrong: {got_r:?}"));
}
if !link_tree.join("F").join("frame_00001.png").is_file() {
return Err("group 5: linked image missing".into());
}
if !masks_link_tree
.join("F")
.join("frame_00001.png.png")
.is_file()
{
return Err("group 5: linked mask missing".into());
}
if masks_link_tree
.join("F")
.join("frame_00002.png.png")
.exists()
{
return Err("group 5: unexpected frame_00002 mask".into());
}
if !masks_link_tree
.join("R")
.join("frame_00003.png.png")
.is_file()
{
return Err("group 5: partial frame_00003 mask missing".into());
}
ok += 1;
// Group 6: partial rig (only 4 faces).
let partial_dir = root.join("partial");
for i in 1..=2u32 {
for face in ["F", "R", "B", "L"] {
write_dummy_png(&partial_dir.join(format!("frame_{i:05}_{face}.png")), n, n)
.map_err(|e| format!("group 6: write partial: {e}"))?;
}
}
let out2 = root.join("rig_partial.json");
let _ = rig::run_rig_stage(
&partial_dir,
&out2,
None,
None,
Layout::Auto,
"F",
"linspace",
None,
None,
LinkMode::Copy,
)
.map_err(|e| format!("group 6: run_rig_stage: {e}"))?;
let text2 = std::fs::read_to_string(&out2).map_err(|e| format!("group 6: read rig: {e}"))?;
let cfg2 = json::parse(&text2).map_err(|e| format!("group 6: parse: {e}"))?;
let cams2 = cfg2
.as_array()
.and_then(|a| a.first())
.and_then(|v| v.find("cameras"))
.and_then(|v| v.as_array())
.ok_or("group 6: no cameras")?;
if cams2.len() != 4 {
return Err(format!("group 6: expected 4 cameras, got {}", cams2.len()));
}
ok += 1;
let _ = std::fs::remove_dir_all(root);
Ok(ok)
}
#[cfg(test)]
mod tests {
use super::run_selftest;
#[test]
fn rig_selftest_passes() {
let ok = run_selftest().expect("self-test must pass");
assert_eq!(ok, 6, "expected 6 check groups");
}
}
+553
View File
@@ -0,0 +1,553 @@
//! Stage 3 — project equirect frames (+masks) to 6 cube faces.
//!
//! GPU compute kernel in `gpu/shaders/equirect_to_cubemap.wgsl`; orchestration
//! here uploads equirect RGBA textures, dispatches per face, reads back the
//! face tiles, and writes `{stem}_{FACE}.ext` (+ `.png` masks) via `image`.
//!
//! Batching: frames are processed in groups of `batch`. All uploads +
//! per-(frame,face) dispatches for a batch go into ONE command encoder, and all
//! face textures are read back with a single `poll(Wait)` — amortising the
//! readback latency and letting the GPU pipeline across frames.
//!
//! Filename / layout parity with the Python `run_cubemap_stage`:
//! * frames: `{stem}_{FACE}{suffix}` (suffix preserved; jpg/jpeg honour
//! `--cubemap-jpeg-quality`, png does not)
//! * masks: `{stem}_{FACE}{suffix}.png` (single-channel binary), only when
//! `masks_dir` (and `output_masks_dir`) are given and a `{stem}.png` mask
//! exists in `masks_dir`.
//! * `--cubemap-skip-faces` and `--cubemap-resume` behave as in Python.
use std::path::{Path, PathBuf};
use anyhow::{anyhow, bail, Context, Result};
use image::{DynamicImage, ImageFormat, ImageReader};
use crate::error::CancelFlag;
use crate::gpu::GpuContext;
/// Face letters in canonical order (matches `math::FACES` and the WGSL `face_id`).
pub const FACES: [&str; 6] = ["F", "R", "B", "L", "U", "D"];
#[repr(C)]
#[derive(Clone, Copy)]
struct Params {
src_w: u32,
src_h: u32,
face_n: u32,
face_id: u32,
mode: u32,
pad0: u32,
pad1: u32,
pad2: u32,
}
fn pod_slice<T: Copy>(x: &T) -> &[u8] {
unsafe { std::slice::from_raw_parts((x as *const T) as *const u8, std::mem::size_of::<T>()) }
}
/// Batched equirect→cubemap GPU renderer. Built once a source frame size is
/// known; resources are pooled and reused across every batch.
pub struct CubemapRenderer<'g> {
gpu: &'g GpuContext,
pipeline: wgpu::ComputePipeline,
#[allow(dead_code)]
bgl: wgpu::BindGroupLayout,
#[allow(dead_code)]
sampler: wgpu::Sampler,
/// One uniform buffer per face_id (0..5), shared across all slots/batches.
/// Re-written at the start of each pass with the current src dims + mode.
uniforms: [wgpu::Buffer; 6],
/// `batch` source equirect textures (re-uploaded per batch).
src_texs: Vec<wgpu::Texture>,
/// `batch * 6` output face storage textures, indexed `slot * 6 + face_id`.
out_texs: Vec<wgpu::Texture>,
/// `batch * 6` bind groups, pre-built once (reference the pools above).
bind_groups: [Vec<wgpu::BindGroup>; 2],
groups: u32,
face_size: u32,
src_w: u32,
src_h: u32,
batch: usize,
}
impl<'g> CubemapRenderer<'g> {
pub fn new(
gpu: &'g GpuContext,
face_size: u32,
src_w: u32,
src_h: u32,
batch: usize,
) -> Result<Self> {
let batch = batch.max(1);
let bgl = gpu.device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("cubemap-bgl"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0, visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform, has_dynamic_offset: false, min_binding_size: None,
}, count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 1, visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: true },
view_dimension: wgpu::TextureViewDimension::D2, multisampled: false,
}, count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 2, visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 3, visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::StorageTexture {
access: wgpu::StorageTextureAccess::WriteOnly,
format: wgpu::TextureFormat::Rgba8Unorm,
view_dimension: wgpu::TextureViewDimension::D2,
}, count: None,
},
],
});
let layout = gpu.device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("cubemap-pipeline-layout"),
bind_group_layouts: &[&bgl],
push_constant_ranges: &[],
});
let pipeline = gpu.compute_pipeline(
"cubemap-pipeline",
crate::gpu::shaders::EQUIRECT_TO_CUBEMAP,
&layout,
);
let sampler = gpu.device.create_sampler(&wgpu::SamplerDescriptor {
label: Some("cubemap-bilinear-sampler"),
address_mode_u: wgpu::AddressMode::ClampToEdge,
address_mode_v: wgpu::AddressMode::ClampToEdge,
address_mode_w: wgpu::AddressMode::ClampToEdge,
mag_filter: wgpu::FilterMode::Linear,
min_filter: wgpu::FilterMode::Linear,
..Default::default()
});
// 6 uniform buffers (one per face_id); written per pass.
let uniforms: [wgpu::Buffer; 6] = core::array::from_fn(|f| {
gpu.create_uniform_buffer(&format!("cubemap-params-{f}"), &[0u8; std::mem::size_of::<Params>()])
});
// Pools.
let mut src_texs = Vec::with_capacity(batch);
let mut out_texs = Vec::with_capacity(batch * 6);
for slot in 0..batch {
src_texs.push(gpu.create_texture_2d(&format!("cubemap-src-{slot}"), src_w, src_h));
for face in 0..6u32 {
out_texs.push(gpu.create_storage_texture_2d(
&format!("cubemap-out-{slot}-{face}"),
face_size,
));
}
}
// Bind groups for both modes (0 = images, 1 = masks). Only the uniform
// contents differ; the bind-group *resources* are identical, so both
// arrays point at the same buffers/textures. We keep two copies so the
// per-pass uniform re-write never aliases a bind group the GPU is still
// reading from a different mode — they are actually the same buffers,
// but written in sequence per pass (images pass fully completes its
// readback before the masks pass writes uniforms), so this is safe.
let build_bgs = |label: &str| -> Vec<wgpu::BindGroup> {
let mut bgs = Vec::with_capacity(batch * 6);
for slot in 0..batch {
let src_view = src_texs[slot].create_view(&wgpu::TextureViewDescriptor::default());
for face in 0..6u32 {
let out_view = out_texs[slot * 6 + face as usize]
.create_view(&wgpu::TextureViewDescriptor::default());
bgs.push(gpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some(&format!("{label}-{slot}-{face}")),
layout: &bgl,
entries: &[
wgpu::BindGroupEntry { binding: 0, resource: uniforms[face as usize].as_entire_binding() },
wgpu::BindGroupEntry { binding: 1, resource: wgpu::BindingResource::TextureView(&src_view) },
wgpu::BindGroupEntry { binding: 2, resource: wgpu::BindingResource::Sampler(&sampler) },
wgpu::BindGroupEntry { binding: 3, resource: wgpu::BindingResource::TextureView(&out_view) },
],
}));
}
}
bgs
};
let bind_groups = [build_bgs("cubemap-bg-img"), build_bgs("cubemap-bg-mask")];
Ok(CubemapRenderer {
gpu, pipeline, bgl, sampler, uniforms, src_texs, out_texs, bind_groups,
groups: face_size.div_ceil(8), face_size, src_w, src_h, batch,
})
}
/// Project a batch of `n` source RGBA frames to the given faces.
/// Returns `n * 6` RGB face planes in order `[slot][face]` (`None` for
/// skipped faces).
#[allow(clippy::needless_range_loop)]
pub fn project_batch(
&self,
frames: &[&[u8]],
active_faces: &[u32],
mode: u32,
) -> Result<Vec<Vec<Option<Vec<u8>>>>> {
let n = frames.len();
debug_assert!(n <= self.batch);
// 1. Write the 6 per-face uniforms for this mode + dims.
for &face_id in active_faces {
let params = Params {
src_w: self.src_w, src_h: self.src_h, face_n: self.face_size,
face_id, mode, pad0: 0, pad1: 0, pad2: 0,
};
self.gpu.queue.write_buffer(&self.uniforms[face_id as usize], 0, pod_slice(&params));
}
// 2. Upload source frames.
for (slot, rgba) in frames.iter().enumerate() {
self.gpu.upload_rgba(&self.src_texs[slot], self.src_w, self.src_h, rgba);
}
// 3. One encoder with n * active_faces dispatches.
let bgs = &self.bind_groups[mode as usize];
let mut encoder = self.gpu.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
{
let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("cubemap-pass"), timestamp_writes: None,
});
pass.set_pipeline(&self.pipeline);
for slot in 0..n {
for &face_id in active_faces {
let bg = &bgs[slot * 6 + face_id as usize];
pass.set_bind_group(0, bg, &[]);
pass.dispatch_workgroups(self.groups, self.groups, 1);
}
}
}
self.gpu.queue.submit(Some(encoder.finish()));
// 4. Read back all n * active_faces output textures with one poll.
let mut tex_refs: Vec<&wgpu::Texture> = Vec::with_capacity(n * active_faces.len());
for slot in 0..n {
for &face_id in active_faces {
tex_refs.push(&self.out_texs[slot * 6 + face_id as usize]);
}
}
let raws = self.gpu.read_textures_rgba(&tex_refs, self.face_size, self.face_size)?;
// 5. De-stride RGBA -> RGB and place into [slot][face] results.
let mut out: Vec<Vec<Option<Vec<u8>>>> = (0..n).map(|_| vec![None; 6]).collect();
let mut idx = 0;
for slot in 0..n {
for &face_id in active_faces {
let rgba = &raws[idx];
idx += 1;
let mut rgb = Vec::with_capacity((self.face_size * self.face_size * 3) as usize);
for px in rgba.as_chunks::<4>().0 {
rgb.extend_from_slice(&px[..3]);
}
out[slot][face_id as usize] = Some(rgb);
}
}
Ok(out)
}
}
/// Parse `--cubemap-skip-faces` (comma-separated, case-insensitive).
pub fn parse_skip_faces(raw: &str) -> Result<Vec<u32>> {
let mut ids = Vec::new();
for tok in raw.split(',') {
let t = tok.trim().to_uppercase();
if t.is_empty() {
continue;
}
let pos = FACES.iter().position(|&f| f == t).ok_or_else(|| {
anyhow!("--cubemap-skip-faces has unrecognised face(s) {t}; valid: {FACES:?}")
})?;
ids.push(pos as u32);
}
Ok(ids)
}
/// Gather supported image files in `dir`, sorted, matching `gather_images`.
pub fn gather_images(dir: &Path) -> Result<Vec<PathBuf>> {
let ext: &[&str] = &["jpg", "jpeg", "png", "bmp", "tif", "tiff", "webp"];
let mut frames: Vec<PathBuf> = std::fs::read_dir(dir)
.with_context(|| format!("cannot read directory {}", dir.display()))?
.filter_map(|e| e.ok())
.map(|e| e.path())
.filter(|p| {
p.is_file()
&& p.extension()
.map(|x| ext.contains(&x.to_string_lossy().to_lowercase().as_ref()))
.unwrap_or(false)
})
.collect();
frames.sort();
if frames.is_empty() {
bail!("No image files found in {}", dir.display());
}
Ok(frames)
}
fn load_rgba(path: &Path) -> Result<(Vec<u8>, u32, u32)> {
let img = ImageReader::open(path)
.with_context(|| format!("cannot open {}", path.display()))?
.with_guessed_format()
.with_context(|| format!("cannot guess format for {}", path.display()))?
.decode()
.with_context(|| format!("cannot decode {}", path.display()))?;
let w = img.width();
let h = img.height();
Ok((img.to_rgba8().into_raw(), w, h))
}
fn load_grayscale_mask(path: &Path) -> Result<Option<(Vec<u8>, u32, u32)>> {
if !path.exists() {
return Ok(None);
}
let img = ImageReader::open(path)
.with_context(|| format!("cannot open mask {}", path.display()))?
.with_guessed_format()?
.decode()
.with_context(|| format!("cannot decode mask {}", path.display()))?;
let w = img.width();
let h = img.height();
Ok(Some((img.to_luma8().into_raw(), w, h)))
}
fn write_face(path: &Path, rgb: &[u8], n: u32, jpeg_quality: u32) -> Result<()> {
let format = match path.extension().and_then(|e| e.to_str()).map(|s| s.to_lowercase()).as_deref() {
Some("jpg") | Some("jpeg") => Some(ImageFormat::Jpeg),
Some("png") => Some(ImageFormat::Png),
Some("bmp") => Some(ImageFormat::Bmp),
Some("webp") => Some(ImageFormat::WebP),
_ => None,
};
let img: DynamicImage = DynamicImage::ImageRgb8(
image::RgbImage::from_raw(n, n, rgb.to_vec())
.ok_or_else(|| anyhow!("bad face dims"))?,
);
match format {
Some(ImageFormat::Jpeg) => {
let mut out = std::fs::File::create(path)
.with_context(|| format!("cannot create {}", path.display()))?;
let mut enc =
image::codecs::jpeg::JpegEncoder::new_with_quality(&mut out, jpeg_quality as u8);
enc.encode_image(&img)
.with_context(|| format!("cannot encode jpeg {}", path.display()))?;
}
Some(fmt) => img.save_with_format(path, fmt)?,
None => img.save(path)?,
}
Ok(())
}
fn write_mask(path: &Path, gray: &[u8], n: u32) -> Result<()> {
let img = image::GrayImage::from_raw(n, n, gray.to_vec())
.ok_or_else(|| anyhow!("bad mask dims"))?;
img.save(path)
.with_context(|| format!("cannot write mask {}", path.display()))?;
Ok(())
}
/// Project every frame in `frames_dir` (and its masks, if present) to faces.
///
/// Returns `(output_frames_dir, output_masks_dir_or_None)`.
#[allow(clippy::too_many_arguments)]
pub fn run_cubemap_stage<'a>(
frames_dir: &Path,
output_frames_dir: &Path,
masks_dir: Option<&Path>,
output_masks_dir: Option<&Path>,
face_size: u32,
skip_faces: &[u32],
jpeg_quality: u32,
resume: bool,
batch: u32,
gpu: &GpuContext,
should_cancel: &CancelFlag,
#[allow(clippy::needless_range_loop)]
logger: crate::logging::Logger<'a>,
) -> Result<(PathBuf, Option<PathBuf>)> {
std::fs::create_dir_all(output_frames_dir)?;
if let Some(omd) = output_masks_dir {
std::fs::create_dir_all(omd)?;
}
let active_faces: Vec<u32> = (0..6).filter(|f| !skip_faces.contains(f)).collect();
let active_names: Vec<&str> = active_faces.iter().map(|&f| FACES[f as usize]).collect();
crate::logging::emit(
&format!(
"cubemap: {} -> {} face_size={} skip=[{}] batch={} (GPU)",
frames_dir.display(),
output_frames_dir.display(),
face_size,
active_names.join(","),
batch,
),
logger,
);
let frames = gather_images(frames_dir)?;
crate::logging::emit(&format!("Found {} equirect frames in {}", frames.len(), frames_dir.display()), logger);
// Probe the first frame's dimensions to size the GPU source textures.
if frames.is_empty() {
bail!("No frames to process.");
}
let (_, w0, h0) = load_rgba(&frames[0])?;
let batch = batch.max(1) as usize;
let renderer = CubemapRenderer::new(gpu, face_size, w0, h0, batch)?;
let total = frames.len() * active_faces.len();
let mut frames_written = 0usize;
let mut masks_written = 0usize;
let mut unreadable = 0usize;
// Buffer decoded frames + per-frame metadata for batched processing.
let mut batch_idx = 0usize;
let mut processed = 0usize;
while batch_idx < frames.len() {
should_cancel.check().map_err(|_| anyhow!("Cancelled."))?;
// Build the batch: load up to `batch` frames (and their masks).
let end = (batch_idx + batch).min(frames.len());
let mut srcs: Vec<Vec<u8>> = Vec::with_capacity(end - batch_idx);
let mut frame_metas: Vec<(String, String)> = Vec::with_capacity(end - batch_idx);
let mut mask_srcs: Vec<Option<Vec<u8>>> = Vec::with_capacity(end - batch_idx);
for (fi, frame_path) in frames[batch_idx..end].iter().enumerate() {
let _fi = batch_idx + fi;
let stem = frame_path.file_stem().map(|s| s.to_string_lossy().into_owned()).unwrap_or_default();
let suffix = frame_path.extension().map(|s| format!(".{}", s.to_string_lossy())).unwrap_or_default();
// Resume: skip a frame only if every active face already exists.
if resume {
let all_exist = active_names.iter().all(|f| {
output_frames_dir.join(format!("{stem}_{f}{suffix}")).exists()
});
if all_exist {
processed += active_faces.len();
crate::logging::emit_progress(&format!("{processed} / {total} (skipped, resume)"), logger);
continue;
}
}
let (rgba, w, h) = match load_rgba(frame_path) {
Ok(x) => x,
Err(e) => {
unreadable += 1;
crate::logging::emit(&format!("[warn] could not read {} ({e})", frame_path.display()), logger);
processed += active_faces.len();
continue;
}
};
if w != w0 || h != h0 {
bail!("frame {} is {w}x{h} but first frame was {w0}x{h0}; mixed dimensions are not supported",
frame_path.display());
}
srcs.push(rgba);
frame_metas.push((stem.clone(), suffix.clone()));
// Optional matching equirect mask. The masks stage writes
// `<frame.name>.png` (e.g. `output_frame_00001.png.png`), matching
// the Python `_cubemap_input_mask_path` / `mask_path_for`.
mask_srcs.push(if let (Some(md), Some(_)) = (masks_dir, output_masks_dir) {
let mask_path = md.join(format!("{stem}{suffix}.png"));
match load_grayscale_mask(&mask_path)? {
Some((gray, mw, mh)) if mw == w && mh == h => {
let mut rgba = Vec::with_capacity((w * h * 4) as usize);
for &g in &gray {
rgba.extend_from_slice(&[g, g, g, 255]);
}
Some(rgba)
}
_ => None,
}
} else {
None
});
}
if !srcs.is_empty() {
let n = srcs.len();
let refs: Vec<&[u8]> = srcs.iter().map(|v| v.as_slice()).collect();
let face_results = renderer.project_batch(&refs, &active_faces, 0)?;
// Masks pass (same pool, mode=1).
let mask_results: Option<Vec<Vec<Option<Vec<u8>>>>> = if output_masks_dir.is_some() && mask_srcs.iter().any(|m| m.is_some()) {
let mrefs: Vec<&[u8]> = mask_srcs.iter().filter_map(|m| m.as_deref()).collect();
// Only frames that actually have a mask are projected; keep an
// alignment map so we know which slot each mask result belongs to.
let mut idx_map = Vec::with_capacity(n);
let mut filtered: Vec<&[u8]> = Vec::new();
for (slot, m) in mask_srcs.iter().enumerate() {
if m.is_some() {
idx_map.push(slot);
filtered.push(m.as_deref().unwrap());
}
}
let _ = mrefs;
if filtered.is_empty() {
None
} else {
let mr = renderer.project_batch(&filtered, &active_faces, 1)?;
// Re-associate results to their original slots.
let mut by_slot: Vec<Vec<Option<Vec<u8>>>> = (0..n).map(|_| vec![None; 6]).collect();
for (k, slot) in idx_map.iter().enumerate() {
by_slot[*slot] = mr[k].clone();
}
Some(by_slot)
}
} else {
None
};
// Write faces + masks for each frame in the batch.
for slot in 0..n {
let (stem, suffix) = &frame_metas[slot];
for &face_id in &active_faces {
let fname = FACES[face_id as usize];
let out_name = format!("{stem}_{fname}{suffix}");
if let Some(rgb) = &face_results[slot][face_id as usize] {
write_face(&output_frames_dir.join(&out_name), rgb, face_size, jpeg_quality)?;
frames_written += 1;
}
if let Some(mr) = &mask_results {
if let Some(mask_rgb) = &mr[slot][face_id as usize] {
let gray: Vec<u8> = mask_rgb.as_chunks::<3>().0.iter()
.map(|px| if px[0] > 127 { 255 } else { 0 }).collect();
write_mask(
&output_masks_dir.unwrap().join(format!("{out_name}.png")),
&gray,
face_size,
)?;
masks_written += 1;
}
}
}
}
}
processed = end * active_faces.len();
crate::logging::emit_progress(&format!("{processed} / {total}"), logger);
batch_idx = end;
}
crate::logging::emit(
&format!("Done. {frames_written} face-crop images written to {}", output_frames_dir.display()),
logger,
);
if let Some(omd) = output_masks_dir {
crate::logging::emit(&format!("{masks_written} face-crop masks written to {}", omd.display()), logger);
}
if unreadable > 0 {
crate::logging::emit(&format!("{unreadable} frames were unreadable and skipped"), logger);
}
Ok((
output_frames_dir.to_path_buf(),
output_masks_dir.map(|p| p.to_path_buf()),
))
}
File diff suppressed because it is too large Load Diff
+991
View File
@@ -0,0 +1,991 @@
//! Stage 3 — YOLO segmentation masks on cubemap face images.
//!
//! Runs *after* the cubemap stage and operates directly on the cube-face images
//! it produced (`03_cubemap/images/{stem}_{FACE}.ext`). This avoids the wasteful
//! round-trip the old design did (e2c -> YOLO -> c2e -> equirect mask, then e2c
//! *again* in the cubemap stage). Cube faces are square and distortion-free, so
//! YOLO detects cleanly and we keep the masks as cube-face masks
//! (`03_cubemap/masks/{stem}_{FACE}.ext.png`) — no reprojection, no stitch-back.
//!
//! Because the inputs are square, the only resize is a uniform scale to the
//! model's `imgsz` (pad 0), so there is no letterbox crop/offset to mis-scale
//! on the way back — the bug that previously produced shifted masks is gone by
//! construction.
//!
//! GPU: ONNX Runtime via the `ort` crate, loaded dynamically from the *system*
//! libonnxruntime (`--features masks` → `ort/load-dynamic`). 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 if the EP is unavailable.
//!
//! Without the `masks` feature this module reports a clear "not enabled" error.
use std::path::{Path, PathBuf};
#[allow(unused_imports)]
use anyhow::{anyhow, bail, Context, Result};
use crate::error::CancelFlag;
use crate::gpu::GpuContext;
#[allow(unused_imports)]
use crate::logging::{self, Logger};
/// COCO 80 class names (lowercased for matching `--mask-classes`).
pub const COCO_CLASSES: [&str; 80] = [
"person",
"bicycle",
"car",
"motorcycle",
"airplane",
"bus",
"train",
"truck",
"boat",
"traffic light",
"fire hydrant",
"stop sign",
"parking meter",
"bench",
"bird",
"cat",
"dog",
"horse",
"sheep",
"cow",
"elephant",
"bear",
"zebra",
"giraffe",
"backpack",
"umbrella",
"handbag",
"tie",
"suitcase",
"frisbee",
"skis",
"snowboard",
"sports ball",
"kite",
"baseball bat",
"baseball glove",
"skateboard",
"surfboard",
"tennis racket",
"bottle",
"wine glass",
"cup",
"fork",
"knife",
"spoon",
"bowl",
"banana",
"apple",
"sandwich",
"orange",
"broccoli",
"carrot",
"hot dog",
"pizza",
"donut",
"cake",
"chair",
"couch",
"potted plant",
"bed",
"dining table",
"toilet",
"tv",
"laptop",
"mouse",
"remote",
"keyboard",
"cell phone",
"microwave",
"oven",
"toaster",
"sink",
"refrigerator",
"book",
"clock",
"vase",
"scissors",
"teddy bear",
"hair drier",
"toothbrush",
];
pub const COCO_MOVABLE_DEFAULTS: &[&str] =
&["person", "car", "bus", "truck", "bicycle", "motorcycle"];
/// Resolve comma-separated class names to COCO ids.
pub fn resolve_classes(raw: &str) -> Result<Vec<u32>> {
let wanted: Vec<String> = raw
.split(',')
.map(|s| s.trim().to_lowercase())
.filter(|s| !s.is_empty())
.collect();
let mut ids = Vec::new();
for name in &wanted {
match COCO_CLASSES.iter().position(|c| *c == name.as_str()) {
Some(i) => ids.push(i as u32),
None => bail!("class `{name}` not in COCO label set"),
}
}
if ids.is_empty() {
bail!("No valid target classes resolved.");
}
Ok(ids)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Projection {
Cubemap,
Equirect,
}
impl Projection {
pub fn parse(s: &str) -> Result<Self> {
match s {
"cubemap" => Ok(Projection::Cubemap),
"equirect" => Ok(Projection::Equirect),
other => Err(anyhow!(
"unknown --mask-projection `{other}` (cubemap|equirect)"
)),
}
}
}
/// Full masks-stage configuration (mirrors the Python `--mask-*` flags).
#[derive(Debug, Clone)]
pub struct MasksConfig {
pub classes: Vec<u32>,
pub model: PathBuf,
pub device: String,
pub conf: f32,
pub imgsz: u32,
pub dilate: u32,
pub projection: Projection,
pub face_size: u32,
pub skip_faces: Vec<u32>,
pub batch: u32,
pub tile_batch: u32,
pub empty_if_none: bool,
pub iou: f32,
/// Optional directory for the MIGraphX compiled-model cache (GPU path).
/// First run compiles+saves; later runs load — avoids repeating the slow
/// one-time kernel compilation. Defaults to `~/.cache/pipeline_360/migraphx`.
pub ep_cache: Option<PathBuf>,
}
impl Default for MasksConfig {
fn default() -> Self {
MasksConfig {
classes: resolve_classes(&COCO_MOVABLE_DEFAULTS.join(",")).unwrap(),
model: PathBuf::from("yolo11x-seg.onnx"),
device: "cpu".into(),
conf: 0.35,
imgsz: 1280,
dilate: 6,
projection: Projection::Cubemap,
face_size: 1024,
skip_faces: Vec::new(),
batch: 8,
tile_batch: 8,
empty_if_none: false,
iou: 0.45,
ep_cache: None,
}
}
}
#[cfg(feature = "masks")]
mod backend {
use super::*;
use image::{GrayImage, ImageReader};
use ndarray::{Array3, Array4, ArrayView2, ArrayView3, Ix4};
use crate::stages::cubemap::FACES;
const NC: usize = 80;
const NM: usize = 32;
struct Det {
cls: u32,
x1: f32,
y1: f32,
x2: f32,
y2: f32, // in imgsz space
coeffs: [f32; NM],
}
struct Yolo {
session: ort::session::Session,
/// Inference size actually accepted by the model (its fixed H/W, or the
/// requested `--mask-imgsz` if the model has dynamic spatial dims).
imgsz: u32,
/// Whether the model accepts a dynamic batch (else we loop batch=1).
batch_dynamic: bool,
/// Human-readable name of the execution provider actually in use.
ep_name: String,
}
/// Build the execution-provider list for `--mask-device`. Returns the EPs to
/// register and a label for logging. GPU EPs are only compiled in when the
/// matching `ort` feature is enabled (migraphx for AMD/ROCm).
///
/// MIGraphX model caching is handled *outside* this function via the
/// `ORT_MIGRAPHX_CACHE_PATH` env var (MIGraphX compiles+saves there on the
/// first run and loads on later runs). The ort `with_save_model`/`with_load_model`
/// options are intentionally NOT used: with this ort + libonnxruntime combo the
/// save path is mis-received and crashes inference, whereas the env-var cache
/// works reliably. See `setup_migraphx_cache`.
fn build_eps(device: &str) -> Result<Vec<ort::ep::ExecutionProviderDispatch>> {
let d = device.trim().to_lowercase();
match d.as_str() {
"rocm" | "migraphx" | "amd" | "gpu" => {
// MIGraphX is AMD's ROCm GPU execution provider. It is the AMD GPU
// path exposed by the system libonnxruntime on this box.
Ok(vec![ort::ep::MIGraphX::default().with_device_id(0).build()])
}
"cuda" | "nvidia" => {
bail!(
"execution provider `{d}` is not enabled in this build (only MIGraphX/ROCm \
for AMD GPUs is compiled in). Use --mask-device rocm, or build with the \
matching ort feature."
)
}
_ => Ok(Vec::new()), // CPU (default ONNX Runtime EP)
}
}
/// Ensure MIGraphX has a compiled-model cache directory. MIGraphX compiles
/// the ONNX graph to GPU kernels on the first run (slow — minutes for a
/// large model on a new arch) and, if `ORT_MIGRAPHX_CACHE_PATH` points at a
/// writable dir, saves the compiled `.mxr` there and loads it on later runs.
///
/// Without a cache dir MIGraphX still tries to save (to an empty path) and
/// *crashes* inference — so this is required for the GPU path to work at all
/// on this build. We pick `~/.cache/pipeline_360/migraphx` by default, or the
/// dir from `--mask-ep-cache`, and never override a value the user set.
fn setup_migraphx_cache(device: &str, override_dir: Option<&Path>, logger: Logger) {
let d = device.trim().to_lowercase();
if !matches!(d.as_str(), "rocm" | "migraphx" | "amd" | "gpu") {
return;
}
const ENV: &str = "ORT_MIGRAPHX_CACHE_PATH";
const ENV2: &str = "ORT_MIGRAPHX_MODEL_CACHE_PATH";
if std::env::var_os(ENV).is_some() && std::env::var_os(ENV2).is_some() {
return; // user already configured the cache.
}
let dir: PathBuf = match override_dir {
Some(p) => p.to_path_buf(),
None => {
let mut base = dirs_or_home();
base.push("pipeline_360");
base.push("migraphx");
base
}
};
if let Err(e) = std::fs::create_dir_all(&dir) {
logging::emit(
&format!(
"[masks] warning: could not create MIGraphX cache dir {}: {e}; \
GPU model will recompile every run.",
dir.display()
),
logger,
);
return;
}
// MIGraphX reads these to locate its compiled-model cache.
if std::env::var_os(ENV).is_none() {
std::env::set_var(ENV, &dir);
}
if std::env::var_os(ENV2).is_none() {
std::env::set_var(ENV2, &dir);
}
logging::emit(
&format!(
"[masks] MIGraphX compiled-model cache: {} (first run compiles+saves; \
later runs load)",
dir.display()
),
logger,
);
}
/// Return a cache directory root: `$XDG_CACHE_HOME` if set, else `~/.cache`.
fn dirs_or_home() -> PathBuf {
if let Some(x) = std::env::var_os("XDG_CACHE_HOME") {
return PathBuf::from(x);
}
if let Some(h) = std::env::var_os("HOME") {
return PathBuf::from(h).join(".cache");
}
PathBuf::from(".cache")
}
impl Yolo {
/// Open the ONNX model, inspect its input/output shapes, validate it is a
/// *segmentation* model with mask protos, and select the requested
/// execution provider (GPU via MIGraphX when `--mask-device rocm`).
/// Falls back to CPU if the GPU EP cannot be registered.
fn open(model: &Path, requested_imgsz: u32, device: &str) -> Result<Self> {
let eps = build_eps(device)?;
let want_gpu = !eps.is_empty();
let make_session =
|eps: &[ort::ep::ExecutionProviderDispatch]| -> Result<ort::session::Session> {
let mut builder = ort::session::Session::builder()
.map_err(|e| anyhow!("ort builder: {e}"))?;
if !eps.is_empty() {
builder = builder
.with_execution_providers(eps)
.map_err(|e| anyhow!("ort execution providers: {e}"))?;
}
builder
.commit_from_file(model)
.map_err(|e| anyhow!("ort load model {}: {e}", model.display()))
};
// Try the requested EP(s); on failure fall back to CPU so a missing
// GPU runtime degrades gracefully instead of aborting the pipeline.
let (session, ep_name) = match make_session(&eps) {
Ok(s) => (
s,
if want_gpu {
device.to_string()
} else {
"cpu".to_string()
},
),
Err(e) if want_gpu => {
eprintln!(
"[masks] requested execution provider `{device}` failed ({e}); \
falling back to CPU. Check that the GPU runtime (ROCm/MIGraphX) \
is installed and ORT_DYLIB_PATH points at a GPU-enabled libonnxruntime."
);
(make_session(&[])?, "cpu (gpu fallback)".to_string())
}
Err(e) => return Err(e),
};
// Validate: segmentation models have two outputs (output0 boxes+coeffs,
// output1 mask protos). Detection-only models have one output.
let out_names: Vec<String> = session
.outputs()
.iter()
.map(|o| o.name().to_string())
.collect();
if !out_names.iter().any(|n| n == "output1") {
bail!(
"model `{}` is not a segmentation model (it has outputs {out_names:?}; \
expected `output0` + `output1` mask protos). Export the *segmentation* \
model, e.g. `yolo export model=yolo11x-seg.pt format=onnx`.",
model.display()
);
}
// Inspect input dims: [batch, 3, height, width]. -1 = dynamic.
let in_outlet = session
.inputs()
.first()
.ok_or_else(|| anyhow!("model has no input tensor"))?;
let (imgsz, batch_dynamic) = match in_outlet.dtype().tensor_shape() {
Some(shape) if shape.len() == 4 => {
let batch_dyn = shape[0] < 0;
let h = shape[2];
let w = shape[3];
let fixed = if h > 0 && w > 0 { Some(h as u32) } else { None };
let imgsz = fixed.unwrap_or(requested_imgsz);
(imgsz, batch_dyn)
}
_ => (requested_imgsz, true),
};
Ok(Yolo {
session,
imgsz,
batch_dynamic,
ep_name,
})
}
fn imgsz(&self) -> u32 {
self.imgsz
}
/// Run inference on `[B,3,imgsz,imgsz]`; return per-image detections
/// (post NMS, in imgsz space) and the proto tensor `[B, NM, mh, mw]`.
/// If the model has a fixed batch dim, runs one image at a time.
fn infer_batch(
&mut self,
inputs: &Array4<f32>,
conf: f32,
iou_thr: f32,
) -> Result<(Vec<Vec<Det>>, Array4<f32>)> {
let b = inputs.shape()[0];
if self.batch_dynamic || b == 1 {
self.infer_once(inputs, conf, iou_thr)
} else {
// Loop batch=1 (model has a fixed batch dim of 1).
let mut all_dets = Vec::with_capacity(b);
let mut protos_acc: Vec<Array3<f32>> = Vec::with_capacity(b);
for bi in 0..b {
let one = inputs.slice(ndarray::s![bi..bi + 1, .., .., ..]).to_owned();
let (mut d, p) = self.infer_once(&one, conf, iou_thr)?;
all_dets.append(&mut d);
// p is [1, NM, mh, mw]; collapse to [NM, mh, mw].
protos_acc.push(
p.slice(ndarray::s![0, .., .., ..])
.to_owned()
.into_dimensionality::<ndarray::Ix3>()
.map_err(|e| anyhow!("proto collapse: {e}"))?,
);
}
// Stack protos back to [B, NM, mh, mw].
let nm = protos_acc[0].shape()[0];
let mh = protos_acc[0].shape()[1];
let mw = protos_acc[0].shape()[2];
let mut protos_b = Array4::zeros((b, nm, mh, mw));
for (bi, p) in protos_acc.into_iter().enumerate() {
protos_b.slice_mut(ndarray::s![bi, .., .., ..]).assign(&p);
}
Ok((all_dets, protos_b))
}
}
fn infer_once(
&mut self,
inputs: &Array4<f32>,
conf: f32,
iou_thr: f32,
) -> Result<(Vec<Vec<Det>>, Array4<f32>)> {
let view = inputs.view();
let outputs = self
.session
.run(ort::inputs![ort::value::TensorRef::from_array_view(view)
.map_err(|e| anyhow!("ort input: {e}"))?])
.map_err(|e| anyhow!("ort run: {e}"))?;
let o0 = outputs["output0"]
.try_extract_array::<f32>()
.map_err(|e| anyhow!("extract output0: {e}"))?;
let o1 = outputs["output1"]
.try_extract_array::<f32>()
.map_err(|e| anyhow!("extract output1: {e}"))?;
let b = o0.shape()[0];
let mut all = Vec::with_capacity(b);
for bi in 0..b {
let preds = o0.slice(ndarray::s![bi, .., ..]);
all.push(decode(&preds, conf, iou_thr));
}
let protos_b = o1
.to_owned()
.into_dimensionality::<Ix4>()
.map_err(|e| anyhow!("protos dim: {e}"))?;
Ok((all, protos_b))
}
}
/// Decode one image's `[C, A]` predictions into kept detections (per-class
/// NMS). Box coords stay in imgsz space.
fn decode(preds: &ArrayView2<f32>, conf: f32, iou_thr: f32) -> Vec<Det> {
let a = preds.shape()[1];
// Candidate anchors with best class > conf.
let mut cands: Vec<(u32, f32, [f32; 4], [f32; NM])> = Vec::new();
for ai in 0..a {
let cx = preds[[0, ai]];
let cy = preds[[1, ai]];
let w = preds[[2, ai]];
let h = preds[[3, ai]];
let mut best_c = 0u32;
let mut best_s = -1.0f32;
for c in 0..NC {
let s = preds[[4 + c, ai]];
if s > best_s {
best_s = s;
best_c = c as u32;
}
}
if best_s > conf {
let mut coeffs = [0f32; NM];
for k in 0..NM {
coeffs[k] = preds[[4 + NC + k, ai]];
}
cands.push((
best_c,
best_s,
[cx - w / 2.0, cy - h / 2.0, cx + w / 2.0, cy + h / 2.0],
coeffs,
));
}
}
// Per-class greedy NMS.
let mut keep: Vec<usize> = Vec::new();
let mut suppressed = vec![false; cands.len()];
for c in 0..NC {
let mut idxs: Vec<usize> = cands
.iter()
.enumerate()
.filter(|(_, d)| d.0 == c as u32)
.map(|(i, _)| i)
.collect();
idxs.sort_by(|a, b| {
cands[*b]
.1
.partial_cmp(&cands[*a].1)
.unwrap_or(std::cmp::Ordering::Equal)
});
for &i in &idxs {
if suppressed[i] {
continue;
}
keep.push(i);
for &j in idxs.iter().filter(|&&j| j != i) {
if suppressed[j] {
continue;
}
if iou(&cands[i].2, &cands[j].2) > iou_thr {
suppressed[j] = true;
}
}
}
}
keep.into_iter()
.map(|i| {
let (cls, _, box_, coeffs) = &cands[i];
Det {
cls: *cls,
x1: box_[0],
y1: box_[1],
x2: box_[2],
y2: box_[3],
coeffs: *coeffs,
}
})
.collect()
}
fn iou(a: &[f32; 4], b: &[f32; 4]) -> f32 {
let xx1 = a[0].max(b[0]);
let yy1 = a[1].max(b[1]);
let xx2 = a[2].min(b[2]);
let yy2 = a[3].min(b[3]);
let w = (xx2 - xx1).max(0.0);
let h = (yy2 - yy1).max(0.0);
let inter = w * h;
let aa = (a[2] - a[0]).max(0.0) * (a[3] - a[1]).max(0.0);
let ab = (b[2] - b[0]).max(0.0) * (b[3] - b[1]).max(0.0);
let union = aa + ab - inter;
if union <= 0.0 {
0.0
} else {
inter / union
}
}
/// Nearest-neighbour resize of an RGB buffer. Used to scale a square cube
/// face to the model's `imgsz` (and only that — no letterbox padding, so
/// there is no crop/offset to undo on the way back).
fn resize_nearest(rgb: &[u8], w: u32, h: u32, nw: u32, nh: u32) -> Vec<u8> {
if nw == w && nh == h {
return rgb.to_vec();
}
let mut out = vec![0u8; (nw * nh * 3) as usize];
for y in 0..nh {
let sy = (y as u64 * h as u64 / nh.max(1) as u64) as usize;
for x in 0..nw {
let sx = (x as u64 * w as u64 / nw.max(1) as u64) as usize;
let si = (sy * w as usize + sx) * 3;
let di = (y as usize * nw as usize + x as usize) * 3;
out[di..di + 3].copy_from_slice(&rgb[si..si + 3]);
}
}
out
}
/// Scale a square `n x n` RGB face to `imgsz x imgsz` (no padding).
fn resize_square(rgb: &[u8], n: u32, imgsz: u32) -> Vec<u8> {
resize_nearest(rgb, n, n, imgsz, imgsz)
}
fn build_input_tensor(scaled: &[&[u8]], imgsz: u32) -> Array4<f32> {
let n = scaled.len();
let mut t = Array4::zeros((n, 3, imgsz as usize, imgsz as usize));
for (bi, lb) in scaled.iter().enumerate() {
for y in 0..imgsz as usize {
for x in 0..imgsz as usize {
let i = (y * imgsz as usize + x) * 3;
t[[bi, 0, y, x]] = lb[i] as f32 / 255.0;
t[[bi, 1, y, x]] = lb[i + 1] as f32 / 255.0;
t[[bi, 2, y, x]] = lb[i + 2] as f32 / 255.0;
}
}
}
t
}
/// Synthesize a binary mask at `out_w x out_h` from detections. Boxes are
/// in imgsz space; `r`, `pad_left`, `pad_top` map imgsz space back to the
/// original (out) space: out_coord = (imgsz_coord - pad) / r.
///
/// For cube-face inputs the scale is uniform and the pads are 0, so this is
/// a pure `out = imgsz / r` mapping with no offset.
#[allow(clippy::too_many_arguments)]
fn synth_mask(
dets: &[Det],
target_ids: &[u32],
protos: ArrayView3<f32>, // [NM, mh, mw]
imgsz: u32,
out_w: u32,
out_h: u32,
r: f32,
pad_left: u32,
pad_top: u32,
dilate: u32,
) -> (Vec<u8>, usize) {
let mh = protos.shape()[1];
let mw = protos.shape()[2];
let mut mask = vec![0u8; (out_w * out_h) as usize];
let mut hits = 0usize;
let p_scale = (mh as f32) / imgsz as f32; // imgsz -> proto grid
for d in dets {
if !target_ids.contains(&d.cls) {
continue;
}
hits += 1;
// Map box to output (original) space via inverse scale.
let bx1 = ((d.x1 - pad_left as f32) / r)
.max(0.0)
.min(out_w as f32 - 1.0);
let by1 = ((d.y1 - pad_top as f32) / r)
.max(0.0)
.min(out_h as f32 - 1.0);
let bx2 = ((d.x2 - pad_left as f32) / r)
.max(0.0)
.min(out_w as f32 - 1.0);
let by2 = ((d.y2 - pad_top as f32) / r)
.max(0.0)
.min(out_h as f32 - 1.0);
let bw = (bx2 - bx1).max(1.0);
let bh = (by2 - by1).max(1.0);
// Proto crop (imgsz space -> proto grid).
let px1 = (d.x1 * p_scale).max(0.0);
let py1 = (d.y1 * p_scale).max(0.0);
let px2 = (d.x2 * p_scale).min(mw as f32 - 1.0);
let py2 = (d.y2 * p_scale).min(mh as f32 - 1.0);
let ix1 = px1.floor() as i32;
let ix2 = (px2.ceil() as i32).max(ix1 + 1);
let iy1 = py1.floor() as i32;
let iy2 = (py2.ceil() as i32).max(iy1 + 1);
for fy in 0..(bh as i32) {
for fx in 0..(bw as i32) {
let t = if bw > 1.0 {
fx as f32 / (bw - 1.0)
} else {
0.0
};
let u = if bh > 1.0 {
fy as f32 / (bh - 1.0)
} else {
0.0
};
let px = ix1 + ((ix2 - ix1) as f32 * t).round() as i32;
let py = iy1 + ((iy2 - iy1) as f32 * u).round() as i32;
let px = px.clamp(0, mw as i32 - 1);
let py = py.clamp(0, mh as i32 - 1);
let mut val = 0.0f32;
for k in 0..NM {
val += d.coeffs[k] * protos[[k, py as usize, px as usize]];
}
let sig = 1.0 / (1.0 + (-val).exp());
if sig > 0.5 {
let ox = (bx1 as i32 + fx).clamp(0, out_w as i32 - 1);
let oy = (by1 as i32 + fy).clamp(0, out_h as i32 - 1);
mask[(oy * out_w as i32 + ox) as usize] = 255;
}
}
}
}
dilate_mask(&mut mask, out_w, out_h, dilate);
(mask, hits)
}
/// Elliptical binary dilation (CPU). For `dilate > 0` only.
fn dilate_mask(mask: &mut [u8], w: u32, h: u32, dilate: u32) {
if dilate == 0 {
return;
}
let r = dilate as i32;
let (w, h) = (w as i32, h as i32);
let src = mask.to_vec();
for y in 0..h {
for x in 0..w {
if mask[(y * w + x) as usize] != 0 {
continue;
}
let mut hit = false;
'outer: for dy in -r..=r {
for dx in -r..=r {
if dx * dx + dy * dy > r * r {
continue;
}
let xx = x + dx;
let yy = y + dy;
if xx < 0 || yy < 0 || xx >= w || yy >= h {
continue;
}
if src[(yy * w + xx) as usize] != 0 {
hit = true;
break 'outer;
}
}
}
if hit {
mask[(y * w + x) as usize] = 255;
}
}
}
}
fn load_rgba(path: &Path) -> Result<(Vec<u8>, u32, u32)> {
let img = ImageReader::open(path)
.with_context(|| format!("cannot open {}", path.display()))?
.with_guessed_format()?
.decode()
.with_context(|| format!("cannot decode {}", path.display()))?;
Ok((img.to_rgba8().into_raw(), img.width(), img.height()))
}
/// Drop the alpha channel from RGBA -> RGB.
fn rgba_to_rgb(rgba: &[u8]) -> Vec<u8> {
let mut rgb = Vec::with_capacity(rgba.len() * 3 / 4);
for px in rgba.as_chunks::<4>().0 {
rgb.extend_from_slice(&px[..3]);
}
rgb
}
/// Parse the cube-face letter out of a `{stem}_{FACE}.ext` filename.
/// Returns the face index (0..6) or None if it isn't a recognised face.
fn face_index(path: &Path) -> Option<u32> {
let stem = path.file_stem()?.to_str()?;
let (_, face) = stem.rsplit_once('_')?;
FACES.iter().position(|f| *f == face).map(|i| i as u32)
}
/// Process cube-face images from `input_dir` (the cubemap stage's image
/// output) and write one binary mask per face to `output_dir`.
///
/// Masks are named `<image_filename>.png` (e.g. `frame_00001_F.png.png`),
/// matching the rig scanner's expectation. Excluded pixels are black.
pub fn run_masks_stage(
input_dir: &Path,
output_dir: &Path,
cfg: &MasksConfig,
_gpu: &GpuContext,
should_cancel: &CancelFlag,
logger: Logger,
) -> Result<()> {
std::fs::create_dir_all(output_dir)?;
let images = crate::stages::cubemap::gather_images(input_dir)?;
logging::emit(
&format!(
"Found {} cube-face images in {}",
images.len(),
input_dir.display()
),
logger,
);
if cfg.projection == Projection::Equirect {
logging::emit(
"Note: --mask-projection equirect is deprecated — masks now always run on \
cube-face images (run the cubemap stage first). Ignoring the setting.",
logger,
);
}
// Filter to the requested faces.
let active: Vec<(PathBuf, u32, String)> = images
.iter()
.filter_map(|p| {
let fid = face_index(p)?;
if cfg.skip_faces.contains(&fid) {
return None;
}
let name = p.file_name()?.to_string_lossy().into_owned();
Some((p.clone(), fid, name))
})
.collect();
if active.is_empty() {
bail!(
"no cube-face images to mask in {} (after applying --mask-skip-faces?)",
input_dir.display()
);
}
if !cfg.model.exists() {
bail!(
"YOLO model not found: {} (export the .pt to ONNX, e.g. \
`yolo export model=yolo11x-seg.pt format=onnx imgsz={}`)",
cfg.model.display(),
cfg.imgsz
);
}
// For the AMD GPU (MIGraphX) path, point MIGraphX at a compiled-model
// cache dir so the (very slow) one-time kernel compilation happens only
// once. Subsequent runs load the cached `.mxr` in seconds. Done via env
// vars because that is the mechanism this MIGraphX build honours.
setup_migraphx_cache(&cfg.device, cfg.ep_cache.as_deref(), logger);
let mut yolo = Yolo::open(&cfg.model, cfg.imgsz, &cfg.device)?;
let imgsz = yolo.imgsz();
if imgsz != cfg.imgsz {
logging::emit(
&format!("Note: model has fixed input {imgsz}x{imgsz}; using that instead of --mask-imgsz {}.",
cfg.imgsz),
logger,
);
}
if !yolo.batch_dynamic {
logging::emit(
"Note: model has a fixed batch dim; running inference one image at a time.",
logger,
);
}
logging::emit(
&format!(
"masks: model={} device={} (ep={}) conf={} imgsz={} dilate={} batch={} skip=[{}]",
cfg.model.display(),
cfg.device,
yolo.ep_name,
cfg.conf,
imgsz,
cfg.dilate,
cfg.batch,
cfg.skip_faces
.iter()
.map(|f| FACES[*f as usize])
.collect::<Vec<_>>()
.join(","),
),
logger,
);
let batch = cfg.batch.max(1) as usize;
let mut written = 0usize;
let mut skipped = 0usize;
let mut total_detections = 0usize;
let mut offset = 0usize;
while offset < active.len() {
should_cancel.check().map_err(|_| anyhow!("Cancelled."))?;
let end = (offset + batch).min(active.len());
let chunk = &active[offset..end];
// Load + square-resize each face to imgsz.
let mut scaled: Vec<Vec<u8>> = Vec::with_capacity(chunk.len());
let mut metas: Vec<(String, u32, u32)> = Vec::with_capacity(chunk.len()); // (out_name, w, h)
let mut any_ok = false;
for (path, _fid, name) in chunk {
let (rgba, w, h) = match load_rgba(path) {
Ok(x) => x,
Err(e) => {
logging::emit(
&format!("[warn] could not read {} ({e})", path.display()),
logger,
);
metas.push((String::new(), 0, 0));
scaled.push(vec![]);
continue;
}
};
let rgb = rgba_to_rgb(&rgba);
let s = resize_square(&rgb, w.min(h), imgsz);
scaled.push(s);
metas.push((name.clone(), w, h));
any_ok = true;
}
if any_ok {
let refs: Vec<&[u8]> = scaled.iter().map(|v| v.as_slice()).collect();
let tensor = build_input_tensor(&refs, imgsz);
let (dets_batch, protos_b) = yolo.infer_batch(&tensor, cfg.conf, cfg.iou)?;
for (k, dets) in dets_batch.iter().enumerate() {
let (name, w, h) = &metas[k];
if w == &0 || h == &0 {
continue; // unreadable image
}
let protos = protos_b.slice(ndarray::s![k, .., .., ..]);
// Square face, uniform scale, no padding.
let r = imgsz as f32 / (*w as f32).max(*h as f32);
let (raw_mask, hits) = synth_mask(
dets,
&cfg.classes,
protos,
imgsz,
*w,
*h,
r,
0,
0,
cfg.dilate,
);
total_detections += hits;
if hits == 0 && !cfg.empty_if_none {
skipped += 1;
continue;
}
// COLMAP convention: black = excluded. Invert so masked
// (movable) objects are black.
let final_mask: Vec<u8> = raw_mask.iter().map(|&v| 255 - v).collect();
let out_path = output_dir.join(format!("{name}.png"));
GrayImage::from_raw(*w, *h, final_mask)
.ok_or_else(|| anyhow!("bad mask dims {w}x{h}"))?
.save(&out_path)
.with_context(|| format!("cannot write {}", out_path.display()))?;
written += 1;
}
}
logging::emit_progress(&format!("masks: {}/{}", end, active.len()), logger);
offset = end;
}
logging::emit(
&format!(
"Done. {written} face masks written to {}, {skipped} face(s) had no target objects \
(unmasked), {total_detections} total target-object instances masked.",
output_dir.display(),
),
logger,
);
logging::emit(&format!("Masks in: {}", output_dir.display()), logger);
Ok(())
}
}
#[cfg(feature = "masks")]
pub fn run_masks_stage(
input_dir: &Path,
output_dir: &Path,
cfg: &MasksConfig,
gpu: &GpuContext,
should_cancel: &CancelFlag,
logger: Logger,
) -> Result<()> {
backend::run_masks_stage(input_dir, output_dir, cfg, gpu, should_cancel, logger)
}
#[cfg(not(feature = "masks"))]
pub fn run_masks_stage(
_input_dir: &Path,
_output_dir: &Path,
_cfg: &MasksConfig,
_gpu: &GpuContext,
_should_cancel: &CancelFlag,
_logger: Logger,
) -> Result<()> {
Err(anyhow!(
"the masks stage requires the `masks` cargo feature (ONNX Runtime). \
Build with: cargo build --features masks"
))
}
+6
View File
@@ -0,0 +1,6 @@
//! Pipeline stage modules, mirroring the Python stage layout.
pub mod cubemap;
pub mod extract;
pub mod masks;
pub mod rig;
+728
View File
@@ -0,0 +1,728 @@
//! Stage 4 — COLMAP rig config + nested `{FACE}/{stem}` tree (pure `std`).
//!
//! Direct port of the Python `run_rig_stage` / `scan_cubemap_*` /
//! `materialise_*` logic. Pure `std` so it compiles and tests offline; produces
//! a `Vec<String>` of log messages that the caller (CLI/GUI) emits.
//!
//! COLMAP matches rig cameras by filename **prefix** only, so the flat
//! `{stem}_{FACE}.ext` stage-3 output must be relinked into a nested
//! `{FACE}/{stem}.ext` tree.
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use crate::json::{self, Value};
use crate::math;
pub type Frames = BTreeMap<String, BTreeMap<String, PathBuf>>;
/// Image file extensions treated as candidates (mirrors Python `IMAGE_EXTENSIONS`).
pub const IMAGE_EXTENSIONS: [&str; 7] = [".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LinkMode {
Symlink,
Hardlink,
Copy,
}
impl LinkMode {
pub fn parse(s: &str) -> Result<LinkMode, String> {
match s {
"symlink" => Ok(LinkMode::Symlink),
"hardlink" => Ok(LinkMode::Hardlink),
"copy" => Ok(LinkMode::Copy),
other => Err(format!("unknown link mode: {other}")),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Layout {
Auto,
Flat,
Nested,
}
fn has_image_suffix(p: &Path) -> bool {
let s = p
.extension()
.and_then(|e| e.to_str())
.unwrap_or("")
.to_lowercase();
IMAGE_EXTENSIONS.iter().any(|e| *e == format!(".{s}"))
}
fn iter_image_files(root: &Path) -> Vec<PathBuf> {
let mut out = Vec::new();
let mut stack = vec![root.to_path_buf()];
while let Some(dir) = stack.pop() {
let Ok(read) = std::fs::read_dir(&dir) else {
continue;
};
for entry in read.flatten() {
let path = entry.path();
if path.is_dir() {
stack.push(path);
} else if has_image_suffix(&path) {
out.push(path);
}
}
}
out.sort();
out
}
/// `{stem}: {FACE}: path` for cubemap images. Supports flat, nested, or auto-detect.
pub fn scan_cubemap_images(images_dir: &Path, layout: Layout) -> Result<(Frames, Layout), String> {
if !images_dir.is_dir() {
return Err(format!(
"--images-dir is not a directory: {}",
images_dir.display()
));
}
let layout = resolve_layout(images_dir, layout)?;
let mut frames: Frames = BTreeMap::new();
match layout {
Layout::Flat => {
for path in iter_image_files(images_dir) {
if let Some((stem, face)) = flat_stem_face(&path) {
frames.entry(stem).or_default().insert(face, path);
}
}
}
_ => {
for face in math::FACES {
let face_dir = images_dir.join(face);
if face_dir.is_dir() {
for path in iter_image_files(&face_dir) {
let stem = path
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("")
.to_string();
frames
.entry(stem)
.or_default()
.insert(face.to_string(), path);
}
}
}
}
}
if frames.is_empty() {
let expected = if layout == Layout::Flat {
"'{stem}_{FACE}.png'"
} else {
"'{FACE}/{stem}.png'"
};
return Err(format!(
"No cubemap face images found under {} (layout={}). Expected {expected}.",
images_dir.display(),
layout_name(layout),
));
}
Ok((frames, layout))
}
/// `{stem}: {FACE}: path` for cubemap masks; a frame/face without a mask is fine.
pub fn scan_cubemap_masks(masks_dir: &Path, layout: Layout) -> Result<(Frames, Layout), String> {
if !masks_dir.is_dir() {
return Err(format!(
"--masks-dir is not a directory: {}",
masks_dir.display()
));
}
let layout = resolve_layout(masks_dir, layout)?;
let mut frames: Frames = BTreeMap::new();
match layout {
Layout::Flat => {
for p in iter_png_files(masks_dir) {
if let Some((stem, face)) = flat_mask_stem_face(&p) {
frames.entry(stem).or_default().insert(face, p);
}
}
}
_ => {
for face in math::FACES {
let face_dir = masks_dir.join(face);
if !face_dir.is_dir() {
continue;
}
for p in iter_png_files(&face_dir) {
// Mask name is "<image_name>.png", i.e. "{stem}.{image_ext}.png".
// file_stem() strips one ext (".png"), leaving "{stem}.{image_ext}";
// drop that image ext to recover the bare stem (works for any image ext).
let stem = p
.file_stem()
.and_then(|s| s.to_str())
.map(|s| s.rsplit_once('.').map(|(h, _)| h).unwrap_or(s))
.unwrap_or("")
.to_string();
frames.entry(stem).or_default().insert(face.to_string(), p);
}
}
}
}
Ok((frames, layout))
}
fn resolve_layout(dir: &Path, layout: Layout) -> Result<Layout, String> {
match layout {
Layout::Auto => {
let subdirs: std::collections::BTreeSet<String> = std::fs::read_dir(dir)
.map(|rd| {
rd.flatten()
.filter(|e| e.path().is_dir())
.filter_map(|e| e.file_name().to_str().map(|s| s.to_uppercase()))
.collect()
})
.unwrap_or_default();
let is_nested =
!subdirs.is_empty() && subdirs.iter().all(|d| math::FACES.contains(&d.as_str()));
Ok(if is_nested {
Layout::Nested
} else {
Layout::Flat
})
}
other => Ok(other),
}
}
fn flat_stem_face(p: &Path) -> Option<(String, String)> {
let stem = p.file_stem()?.to_str()?;
let (head, face) = stem.rsplit_once('_')?;
if math::FACES.contains(&face) {
Some((head.to_string(), face.to_string()))
} else {
None
}
}
fn flat_mask_stem_face(p: &Path) -> Option<(String, String)> {
// Mask naming: "{stem}_{FACE}.{image_ext}.png" (the image filename with
// ".png" appended). Strip the trailing ".png", drop the image extension,
// then split off the face letter. Handles any image extension (.png/.jpg/…).
let name = p.file_name()?.to_str()?;
let base = name.strip_suffix(".png")?; // "frame_00001_F.jpg" (or .png)
let no_ext = base.rsplit_once('.').map(|(h, _)| h).unwrap_or(base); // "frame_00001_F"
let (head, face) = no_ext.rsplit_once('_')?; // ("frame_00001", "F")
if math::FACES.contains(&face) {
Some((head.to_string(), face.to_string()))
} else {
None
}
}
fn iter_png_files(dir: &Path) -> Vec<PathBuf> {
let mut out = Vec::new();
let mut stack = vec![dir.to_path_buf()];
while let Some(d) = stack.pop() {
let Ok(read) = std::fs::read_dir(&d) else {
continue;
};
for entry in read.flatten() {
let path = entry.path();
if path.is_dir() {
stack.push(path);
} else if path
.extension()
.and_then(|e| e.to_str())
.map(|e| e.eq_ignore_ascii_case("png"))
== Some(true)
{
out.push(path);
}
}
}
out.sort();
out
}
/// Read PNG dimensions from the header without decoding (pure std).
pub fn png_dimensions(path: &Path) -> Option<(u32, u32)> {
let mut head = [0u8; 24];
let mut f = std::fs::File::open(path).ok()?;
use std::io::Read;
f.read_exact(&mut head).ok()?;
if &head[..8] != b"\x89PNG\r\n\x1a\n" || &head[12..16] != b"IHDR" {
return None;
}
let w = u32::from_be_bytes([head[16], head[17], head[18], head[19]]);
let h = u32::from_be_bytes([head[20], head[21], head[22], head[23]]);
Some((w, h))
}
/// Detect a consistent square face size from sampled images.
pub fn detect_face_size(frames: &Frames, sample_limit: usize) -> Result<u32, String> {
let mut seen: BTreeMap<(u32, u32), PathBuf> = BTreeMap::new();
let mut checked = 0usize;
'outer: for faces in frames.values() {
for path in faces.values() {
if let Some(dims) = png_dimensions(path) {
seen.entry(dims).or_insert_with(|| path.clone());
}
checked += 1;
if checked >= sample_limit {
break 'outer;
}
}
if checked >= sample_limit {
break;
}
}
if seen.is_empty() {
return Err("could not read any image dimensions (non-PNG without Pillow?). Pass --rig-face-size explicitly.".into());
}
if seen.len() > 1 {
let detail: Vec<String> = seen
.iter()
.map(|((w, h), p)| {
format!(
"{w}x{h} (e.g. {})",
p.file_name().unwrap_or_default().to_string_lossy()
)
})
.collect();
return Err(format!(
"inconsistent face dimensions: {}. Pass --rig-face-size explicitly.",
detail.join(", ")
));
}
let (&(w, h), example) = seen.iter().next().unwrap();
if w != h {
return Err(format!(
"faces must be square, got {w}x{h} ({}).",
example.file_name().unwrap_or_default().to_string_lossy()
));
}
Ok(w)
}
fn present_faces(frames: &Frames) -> Vec<String> {
let set: std::collections::BTreeSet<String> = frames
.iter()
.flat_map(|(_, fs)| fs.keys().cloned())
.collect();
math::FACES
.iter()
.filter(|f| set.contains(**f))
.map(|f| f.to_string())
.collect()
}
/// Build the COLMAP rig config: one rig definition reused for every frame.
///
/// Schema (COLMAP >= 3.12 `rig_configurator`):
/// `[{"cameras": [ {"image_prefix": "F/", "ref_sensor": true, "camera_model_name":
/// "PINHOLE", "camera_params": [fx,fy,cx,cy]}, ... ]}]`
pub fn build_colmap_rig_config(
present_faces: &[String],
ref_face: &str,
intrinsics: (f64, f64, f64, f64),
) -> Result<Value, String> {
let (fx, fy, cx, cy) = intrinsics;
if !present_faces.iter().any(|f| f == ref_face) {
return Err(format!(
"reference face {ref_face:?} not among present faces {present_faces:?}"
));
}
let mut cameras = Vec::new();
let ref_string = ref_face.to_string();
let ordered: Vec<&String> = std::iter::once(&ref_string)
.chain(present_faces.iter().filter(|f| *f != ref_face))
.collect();
for face in ordered {
let mut entry: Vec<(String, Value)> =
vec![("image_prefix".into(), Value::String(format!("{face}/")))];
if face == ref_face {
entry.push(("ref_sensor".into(), Value::Bool(true)));
} else {
let (qw, qx, qy, qz) = math::cam_from_rig_rotation(face, ref_face);
entry.push((
"cam_from_rig_rotation".into(),
Value::Array(vec![
Value::Number(qw),
Value::Number(qx),
Value::Number(qy),
Value::Number(qz),
]),
));
entry.push((
"cam_from_rig_translation".into(),
Value::Array(vec![
Value::Number(0.0),
Value::Number(0.0),
Value::Number(0.0),
]),
));
}
entry.push(("camera_model_name".into(), Value::String("PINHOLE".into())));
entry.push((
"camera_params".into(),
Value::Array(vec![
Value::Number(fx),
Value::Number(fy),
Value::Number(cx),
Value::Number(cy),
]),
));
cameras.push(Value::Object(entry));
}
Ok(Value::Array(vec![Value::Object(vec![(
"cameras".into(),
Value::Array(cameras),
)])]))
}
/// Write the rig config JSON (indented, trailing newline) to `output`.
pub fn write_rig_config(config: &Value, output: &Path) -> Result<(), String> {
let dir = output.parent().unwrap_or_else(|| Path::new("."));
std::fs::create_dir_all(dir).map_err(|e| format!("mkdir {}: {e}", dir.display()))?;
let text = pretty_json(config);
std::fs::write(output, text + "\n").map_err(|e| format!("write {}: {e}", output.display()))
}
/// Simple two-space-indent pretty printer for the rig schema.
fn pretty_json(v: &Value) -> String {
let mut out = String::new();
pretty_write(v, &mut out, 0);
out
}
fn pretty_write(v: &Value, out: &mut String, depth: usize) {
let indent = " ".repeat(depth);
match v {
Value::Array(items) => {
if items.is_empty() {
out.push_str("[]");
return;
}
out.push_str("[\n");
for (i, item) in items.iter().enumerate() {
out.push_str(&" ".repeat(depth + 1));
pretty_write(item, out, depth + 1);
if i + 1 < items.len() {
out.push(',');
}
out.push('\n');
}
out.push_str(&indent);
out.push(']');
}
Value::Object(pairs) => {
if pairs.is_empty() {
out.push_str("{}");
return;
}
out.push_str("{\n");
for (i, (k, val)) in pairs.iter().enumerate() {
out.push_str(&" ".repeat(depth + 1));
json::Value::String(k.clone()).write(out);
out.push_str(": ");
pretty_write(val, out, depth + 1);
if i + 1 < pairs.len() {
out.push(',');
}
out.push('\n');
}
out.push_str(&indent);
out.push('}');
}
other => other.write(out),
}
let _ = indent;
}
/// Relink "flat '{stem}_{FACE}.ext'" as "'{FACE}/{stem}.ext'" under `dest`.
pub fn materialise_nested_tree(
frames: &Frames,
dest: &Path,
mode: LinkMode,
) -> Result<usize, String> {
let mut count = 0usize;
for (stem, faces) in frames {
for (face, src) in faces {
let face_dir = dest.join(face);
std::fs::create_dir_all(&face_dir)
.map_err(|e| format!("mkdir {}: {e}", face_dir.display()))?;
let ext = src.extension().and_then(|e| e.to_str()).unwrap_or("");
let target = face_dir.join(format!("{stem}.{ext}"));
link(src, &target, mode)?;
count += 1;
}
}
Ok(count)
}
/// Relink masks to "'{FACE}/{stem}{image_ext}.png'", mirroring the image tree.
/// Returns `(linked, orphans)`.
pub fn materialise_mask_tree(
mask_frames: &Frames,
image_frames: &Frames,
dest: &Path,
mode: LinkMode,
) -> Result<(usize, usize), String> {
let mut linked = 0usize;
let mut orphans = 0usize;
for (stem, faces) in mask_frames {
for (face, mask_src) in faces {
let image_src = image_frames.get(stem).and_then(|fs| fs.get(face));
let mask_name = match image_src {
Some(img) => {
let ext = img.extension().and_then(|e| e.to_str()).unwrap_or("");
format!("{stem}.{ext}.png")
}
None => {
orphans += 1;
format!("{stem}.png.png")
}
};
let face_dir = dest.join(face);
std::fs::create_dir_all(&face_dir)
.map_err(|e| format!("mkdir {}: {e}", face_dir.display()))?;
link(mask_src, &face_dir.join(mask_name), mode)?;
linked += 1;
}
}
Ok((linked, orphans))
}
fn link(src: &Path, target: &Path, mode: LinkMode) -> Result<(), String> {
if target.exists() || target.is_symlink() {
let _ = std::fs::remove_file(target);
}
let resolved = src.canonicalize().unwrap_or_else(|_| src.to_path_buf());
match mode {
LinkMode::Symlink => std::os::unix::fs::symlink(&resolved, target).map_err(|e| {
format!(
"symlink {} -> {}: {e}",
target.display(),
resolved.display()
)
}),
LinkMode::Hardlink => {
std::fs::hard_link(&resolved, target).map_err(|e| format!("hardlink {target:?}: {e}"))
}
LinkMode::Copy => std::fs::copy(&resolved, target)
.map(|_| ())
.map_err(|e| format!("copy {target:?}: {e}")),
}
}
/// Result of running stage 4.
#[derive(Debug)]
pub struct RigResult {
pub final_images_path: PathBuf,
pub final_masks_path: Option<PathBuf>,
pub messages: Vec<String>,
}
/// Run the full stage-4 rig generation. `messages` carries human-readable log lines.
pub fn run_rig_stage(
images_dir: &Path,
output: &Path,
masks_dir: Option<&Path>,
face_size: Option<u32>,
layout: Layout,
ref_face: &str,
intrinsics_convention: &str,
link_tree: Option<&Path>,
masks_link_tree: Option<&Path>,
link_mode: LinkMode,
) -> Result<RigResult, String> {
if masks_link_tree.is_some() && masks_dir.is_none() {
return Err("masks_link_tree requires masks_dir".into());
}
let mut m = Vec::new();
let (frames, resolved_layout) = scan_cubemap_images(images_dir, layout)?;
let faces = present_faces(&frames);
let mut ref_face = ref_face.to_string();
if !faces.contains(&ref_face) {
m.push(format!(
"warning: --rig-ref-face {ref_face} not present; using {} as rig frame.",
faces[0]
));
ref_face = faces[0].clone();
}
let missing_faces: Vec<&str> = math::FACES
.iter()
.filter(|f| !faces.iter().any(|s| s.as_str() == **f))
.copied()
.collect();
let resolved_face_size = match face_size {
Some(s) => s,
None => detect_face_size(&frames, 24)?,
};
let intrinsics = math::compute_intrinsics(resolved_face_size, intrinsics_convention);
let (fx, fy, cx, cy) = intrinsics;
let config = build_colmap_rig_config(&faces, &ref_face, intrinsics)?;
write_rig_config(&config, output)?;
let partial: Vec<(String, Vec<String>)> = frames
.iter()
.filter(|(_, fs)| {
fs.keys().collect::<std::collections::BTreeSet<_>>() != faces.iter().collect()
})
.map(|(s, fs)| (s.clone(), fs.keys().cloned().collect()))
.collect();
m.push(format!(
"scanned : {} (layout={})",
images_dir.display(),
layout_name(resolved_layout)
));
m.push(format!("frames : {}", frames.len()));
m.push(format!(
"faces present: {} absent: {}",
faces.join(","),
if missing_faces.is_empty() {
"-".to_string()
} else {
missing_faces.join(",")
}
));
m.push(format!(
"face size : {}px{}",
resolved_face_size,
if face_size.is_some() {
" (override)"
} else {
" (auto-detected)"
}
));
m.push(format!(
"intrinsics : PINHOLE fx={fx:.6} fy={fy:.6} cx={cx:.6} cy={cy:.6} [{intrinsics_convention}]",
));
m.push(format!(
"rig frame : {ref_face} ({}) -- 1 rig definition, {} frames, {} cameras",
face_name(&ref_face),
frames.len(),
faces.len()
));
if !partial.is_empty() {
let shown: Vec<String> = partial
.iter()
.take(3)
.map(|(s, fs)| format!("{s}=[{}]", fs.join(",")))
.collect();
m.push(format!(
"partial frames: {} frame(s) missing >=1 face, e.g. {}",
partial.len(),
shown.join("; ")
));
}
m.push(format!("wrote : {}", output.display()));
let mut final_images = images_dir.to_path_buf();
if let Some(lt) = link_tree {
let n = materialise_nested_tree(&frames, lt, link_mode)?;
m.push(format!(
"linked : {} images into {} ({})",
n,
lt.display(),
link_mode_name(link_mode)
));
final_images = lt.to_path_buf();
} else if resolved_layout == Layout::Flat {
m.push("".into());
m.push(
"IMPORTANT: COLMAP matches rig cameras by filename PREFIX only, so this flat".into(),
);
m.push(
" '{stem}_{FACE}.png' naming cannot be used directly with --rig-link-tree unset."
.into(),
);
}
let mut final_masks: Option<PathBuf> = None;
if let Some(md) = masks_dir {
let (mask_frames, mask_layout) = scan_cubemap_masks(md, Layout::Auto)?;
let n_masks: usize = mask_frames.values().map(|f| f.len()).sum();
let n_images: usize = frames.values().map(|f| f.len()).sum();
m.push("".into());
m.push(format!(
"masks scanned: {} (layout={})",
md.display(),
layout_name(mask_layout)
));
m.push(format!(
"masks found : {n_masks} across {} frame(s) (of {n_images} face images -- missing ones are unmasked)",
mask_frames.len()
));
let unknown: Vec<&String> = mask_frames
.keys()
.filter(|s| !frames.contains_key(*s))
.collect();
if !unknown.is_empty() {
m.push(format!(
"warning: {} mask frame(s) have no matching image frame (stem drift?), e.g. {}",
unknown.len(),
unknown
.iter()
.take(3)
.map(|s| s.as_str())
.collect::<Vec<_>>()
.join(", ")
));
}
if let Some(mlt) = masks_link_tree {
let (linked, orphans) = materialise_mask_tree(&mask_frames, &frames, mlt, link_mode)?;
m.push(format!(
"masks linked : {linked} into {} ({})",
mlt.display(),
link_mode_name(link_mode)
));
if orphans > 0 {
m.push(format!(
"warning: {orphans} mask(s) linked with a fallback name, no matching image found -- verify by hand."
));
}
final_masks = Some(mlt.to_path_buf());
} else if mask_layout == Layout::Flat {
m.push("".into());
m.push("IMPORTANT: masks must mirror the nested image tree too. Set --rig-masks-link-tree.".into());
final_masks = Some(md.to_path_buf());
}
}
Ok(RigResult {
final_images_path: final_images,
final_masks_path: final_masks,
messages: m,
})
}
fn face_name(face: &str) -> &str {
math::FACE_NAMES
.iter()
.find(|(f, _)| *f == face)
.map(|(_, n)| *n)
.unwrap_or(face)
}
fn layout_name(l: Layout) -> &'static str {
match l {
Layout::Flat => "flat",
Layout::Nested => "nested",
Layout::Auto => "auto",
}
}
fn link_mode_name(m: LinkMode) -> &'static str {
match m {
LinkMode::Symlink => "symlink",
LinkMode::Hardlink => "hardlink",
LinkMode::Copy => "copy",
}
}
+154
View File
@@ -0,0 +1,154 @@
//! Video decode (stage 0).
//!
//! Decoding stays on the CPU via FFmpeg, but instead of fragile native
//! bindings (`ffmpeg-next`) we shell out to the **`ffmpeg` / `ffprobe` CLI**
//! and pipe raw `rgba` frames. This is robust across FFmpeg versions (the
//! system here ships FFmpeg 9 / libavcodec 63, which `ffmpeg-next 7` would not
//! match), needs no system dev headers, and is plenty fast for sequential
//! frame iteration.
//!
//! Two passes mirror the Python design:
//! * **analysis** — decode downscaled (to `scale_width`) RGBA, hand to the
//! GPU sharpness/quality kernels;
//! * **extraction** — decode full-resolution RGBA, write the winning frames.
use std::io::Read;
use std::path::Path;
use std::process::{Command, Stdio};
use anyhow::{anyhow, Context, Result};
/// Probed video geometry + (best-effort) frame count.
#[derive(Debug, Clone)]
pub struct VideoProbe {
pub width: u32,
pub height: u32,
/// Declared frame count if the container reports one.
pub nb_frames: Option<u64>,
/// Nominal fps (numerator/denominator), e.g. (30000, 1001).
pub fps: Option<(u64, u64)>,
}
impl VideoProbe {
/// Nominal fps as an f64, if known.
pub fn fps_f64(&self) -> Option<f64> {
self.fps.map(|(n, d)| n as f64 / d as f64)
}
}
/// Probe a video with `ffprobe` (json).
pub fn probe(path: &Path) -> Result<VideoProbe> {
let output = Command::new("ffprobe")
.arg("-v").arg("error")
.arg("-select_streams").arg("v:0")
.arg("-show_entries").arg("stream=width,height,nb_frames,r_frame_rate")
.arg("-of").arg("json")
.arg(path)
.output()
.context("failed to spawn `ffprobe` (is ffmpeg installed?)")?;
if !output.status.success() {
return Err(anyhow!("ffprobe failed: {}", String::from_utf8_lossy(&output.stderr)));
}
let json = String::from_utf8_lossy(&output.stdout);
let width = parse_json_field(&json, "width").context("no width in ffprobe output")?;
let height = parse_json_field(&json, "height").context("no height in ffprobe output")?;
let nb_frames = parse_json_field(&json, "nb_frames").ok().map(|n| n as u64);
let fps = parse_json_string(&json, "r_frame_rate").and_then(|s| {
let (n, d) = s.split_once('/')?;
Some((n.parse().ok()?, d.parse().ok()?))
});
Ok(VideoProbe { width, height, nb_frames, fps })
}
fn parse_json_field(json: &str, key: &str) -> Result<u32> {
let s = parse_json_string(json, key).ok_or_else(|| anyhow!("missing `{key}`"))?;
s.parse::<u32>().map_err(|_| anyhow!("`{key}` not an integer: `{s}`"))
}
/// Tiny regex-free JSON string extractor for the flat `ffprobe -of json` shape.
fn parse_json_string(json: &str, key: &str) -> Option<String> {
let needle = format!("\"{key}\"");
let i = json.find(&needle)?;
let rest = &json[i + needle.len()..];
let colon = rest.find(':')?;
let rest = &rest[colon + 1..];
let rest = rest.trim_start();
if let Some(inner) = rest.strip_prefix('"') {
// Quoted string value.
let end = inner.find('"')?;
return Some(inner[..end].to_string());
}
// Bare number / token up to comma/brace/newline.
let end = rest.find(|c: char| c == ',' || c == '}' || c.is_whitespace()).unwrap_or(rest.len());
Some(rest[..end].trim().to_string())
}
/// Output dimensions after an optional width-preserving downscale.
pub fn scaled_dims(probe: &VideoProbe, scale_width: u32) -> (u32, u32) {
if scale_width == 0 || probe.width <= scale_width {
(probe.width, probe.height)
} else {
let h = ((u64::from(probe.height) * u64::from(scale_width))
.div_ceil(u64::from(probe.width))) as u32;
(scale_width, h.max(1))
}
}
/// A sequential raw-RGBA frame reader piped from `ffmpeg`.
pub struct FramePipe {
child: std::process::Child,
stdout: std::process::ChildStdout,
pub width: u32,
pub height: u32,
frame_bytes: usize,
}
impl FramePipe {
/// Open `path` and decode frames sequentially.
///
/// `scale_width > 0` downscales preserving aspect (analysis pass);
/// `scale_width == 0` decodes at native resolution (extraction pass).
pub fn open(path: &Path, scale_width: u32, probe: &VideoProbe) -> Result<Self> {
let (w, h) = scaled_dims(probe, scale_width);
let mut cmd = Command::new("ffmpeg");
cmd.arg("-hide_banner").arg("-loglevel").arg("error")
.arg("-i").arg(path)
.arg("-f").arg("rawvideo").arg("-pix_fmt").arg("rgba");
if scale_width > 0 && probe.width > scale_width {
cmd.arg("-vf").arg(format!("scale={w}:-1"));
}
cmd.arg("pipe:1")
.stdin(Stdio::null()).stderr(Stdio::piped()).stdout(Stdio::piped());
let mut child = cmd.spawn().context("failed to spawn `ffmpeg` (is ffmpeg installed?)")?;
let stdout = child.stdout.take().ok_or_else(|| anyhow!("ffmpeg: no stdout pipe"))?;
Ok(FramePipe { child, stdout, width: w, height: h, frame_bytes: (w as usize) * (h as usize) * 4 })
}
/// Read the next frame's RGBA bytes, or `None` at EOF.
pub fn next_frame(&mut self) -> Result<Option<Vec<u8>>> {
let mut buf = vec![0u8; self.frame_bytes];
let mut filled = 0usize;
while filled < buf.len() {
let n = self.stdout.read(&mut buf[filled..])?;
if n == 0 { break; }
filled += n;
}
if filled == 0 || filled < buf.len() {
return Ok(None);
}
Ok(Some(buf))
}
/// Drop the decoder (kills the child if still running).
pub fn close(mut self) {
let _ = self.child.kill();
let _ = self.child.wait();
}
}
impl Drop for FramePipe {
fn drop(&mut self) {
let _ = self.child.kill();
let _ = self.child.wait();
}
}
+65
View File
@@ -0,0 +1,65 @@
# Architecture
## Source layout
```
Cargo.toml binary + lib; features: default=["masks"], gui
src/
main.rs CLI entry, clap parsing, `all` orchestration, GUI launch
lib.rs crate root, re-exports
project.rs project layout (01_frames … 04_colmap) + file gathering
logging.rs emit / emit_progress ([progress] line protocol)
error.rs error + cancellation types
math.rs face rotations, quaternions, intrinsics (pure std)
json.rs minimal JSON writer (pure std)
selftest.rs rig self-test (6 check groups)
gpu/
mod.rs GpuContext (device/queue), texture/readback helpers
shaders/ equirect_to_cubemap.wgsl, sharpness.wgsl, c2e.wgsl
stages/
extract.rs stage 0: GPU sharpness + selection + frame writes
masks.rs stage 1: YOLO segmentation (ONNX Runtime / MIGraphX)
cubemap.rs stage 2: GPU equirect → cubemap
rig.rs stage 3: rig config + nested tree
video/
mod.rs ffmpeg/ffprobe CLI decode → raw RGBA frames
gui.rs egui+wgpu GUI (feature `gui`)
offline_check/ pure-std fallback build of the rig core (no cargo/network)
```
## How a run flows
1. **CLI** parses the command; `all` builds a `GpuContext` once and runs each
stage in order, sharing the device/queue.
2. **Extract** pipes decoded frames from ffmpeg to the GPU `sharpness` kernel,
reads back tiny per-workgroup metrics, filters by quality, selects frames,
and writes them out.
3. **Cubemap** uploads each equirect frame and dispatches the
`equirect_to_cubemap` kernel once per face, reading back square faces.
4. **Masks** loads `yolo11x-seg.onnx` into ONNX Runtime (with the MIGraphX EP
when requested), runs each cube face through the network, and synthesizes a
binary mask (see [GPU](GPU.md)).
5. **Rig** scans images + masks, writes `rig_config.json`, and links the nested
COLMAP tree (symlink/hardlink/copy).
## Why Rust + WebGPU
- **Parallelism**: per-pixel/per-frame work maps naturally onto a compute
shader; `rayon` covers CPU-side orchestration.
- **Cross-vendor GPU**: one WGSL codebase → Vulkan / Metal / DX12 / WebGPU,
instead of per-vendor libraries.
- **Safety**: Rust's ownership model + wgpu's resource lifetimes avoid
use-after-free and data races.
- **Distribution**: a single binary; no Python environment.
Video **decode** stays on the CPU through ffmpeg because codecs aren't
something you hand-roll on a bindless GPU layer; everything after decode is GPU.
## Feature flags
| Feature | Effect |
| --- | --- |
| `masks` (default) | enable ONNX Runtime YOLO masks (`ort` + `ndarray`, `load-dynamic`, `migraphx`) |
| `gui` | build the native egui+wgpu GUI into the same binary |
`cargo build --no-default-features` gives a light build with no ONNX Runtime.
+75
View File
@@ -0,0 +1,75 @@
# GPU
This project uses **WebGPU** via the `wgpu` crate, so the same code runs on
Vulkan, Metal, DX12, and WebGPU/wasm.
## GPU selection (all stages)
By default the pipeline asks wgpu for the **discrete** (high-performance)
adapter and logs which one it picked:
```
[gpu] using AMD Radeon RX 9070 XT (RADV GFX1201) (DiscreteGpu) via Vulkan
```
If you see an integrated GPU or a software (llvmpipe) adapter, force the right
backend/device. The `request_adapter` fallback enumerates adapters and prefers
a discrete GPU. On dual-GPU laptops set `WGPU_POWER_PREF` / run under the
discrete-GPU driver profile if wgpu still picks the iGPU.
## YOLO masks on the GPU (ONNX Runtime → MIGraphX)
The masks stage runs YOLO through **ONNX Runtime**, loaded **dynamically from
the system libonnxruntime** (`ORT_DYLIB_PATH`) rather than a downloaded
CPU-only build. On AMD/ROCm systems that system lib is built with the
**MIGraphX** execution provider, so:
```bash
# GPU inference on an AMD GPU
pipeline_360 masks --input-dir 03_cubemap/images --output-dir 03_cubemap/masks \
--mask-model yolo11x-seg.onnx --mask-device rocm --mask-imgsz 1280
# CPU (no ROCm/MIGraphX runtime needed)
pipeline_360 masks --input-dir 03_cubemap/images --output-dir 03_cubemap/masks \
--mask-model yolo11x-seg.onnx --mask-device cpu
```
- If the GPU EP can't be registered, the stage **falls back to CPU** and warns.
- The binary is built with `--features masks` by default.
- On your system the GPU-enabled `libonnxruntime` must be findable — the build
uses `load-dynamic`, so it `dlopen`s `libonnxruntime.so` (or whatever
`ORT_DYLIB_PATH` points at) at runtime, and ROCm must be installed.
### Compiled-model cache (important for first run)
MIGraphX compiles the ONNX graph to GPU kernels on the **first run**, which is
slow (minutes for a large model like `yolo11x-seg` on a new GPU). It caches the
compiled `.mxr` and loads it on later runs.
The pipeline **auto-sets `ORT_MIGRAPHX_CACHE_PATH`** to
`~/.cache/pipeline_360/migraphx` (override with `--mask-ep-cache <dir>`), so:
- first run with `--mask-device rocm` → compiles + saves (slow),
- every later run → loads the cached model in seconds.
### Env vars (useful)
| Variable | Meaning |
| --- | --- |
| `ORT_DYLIB_PATH` | path to the libonnxruntime to dlopen at runtime |
| `ORT_MIGRAPHX_CACHE_PATH` | MIGraphX compiled-model cache dir (auto-set) |
| `ORT_MIGRAPHX_MODEL_CACHE_PATH` | alias some MIGraphX builds read |
## Troubleshooting
- **`[gpu] using llvmpipe / iGPU`** — set `WGPU_POWER_PREF=high` or switch to
the discrete GPU.
- **"not a segmentation model"** — you exported the detection model; export the
**seg** one (`yolo11x-seg.pt`), not `yolo11x.onnx`.
- **MIGraphX crash at inference with an empty `.mxr` path** — the cache env var
was unset. The pipeline sets it automatically; if you run via an external
launcher that strips env vars, set `ORT_MIGRAPHX_CACHE_PATH` yourself.
- **First run "very slow"** — that's the one-time MIGraphX kernel compilation.
Let it finish; later runs are fast.
- **CPU-only ONNX Runtime** — if you overrode `ORT_DYLIB_PATH` to a CPU build,
the MIGraphX EP won't be available and it falls back to CPU.
+53
View File
@@ -0,0 +1,53 @@
# Getting Started
## Prerequisites
- A **Rust toolchain** (`cargo`, `rustc`).
- **`ffmpeg` / `ffprobe`** on `PATH` — video decode shells out to the CLI and
pipes raw frames; no native FFmpeg bindings or dev headers are needed.
- A GPU with a **WebGPU backend** — Vulkan (Linux/Windows), Metal (macOS),
DX12 (Windows), or WebGPU/wasm. On Linux, RADV (Mesa) or AMDVLK works.
- For the **masks** stage: an **ONNX** export of a YOLO segmentation model
(e.g. `yolo11x-seg.onnx`). The `.pt` is not loadable in Rust:
```bash
yolo export model=yolo11x-seg.pt format=onnx imgsz=1280 dynamic=True
```
## Build
Masks are enabled by default, so a plain build is the full pipeline:
```bash
cargo build
```
Lighter build without ONNX Runtime (no masks stage):
```bash
cargo build --no-default-features
```
Build with the native GUI:
```bash
cargo build --features gui
```
## First run (sanity check)
```bash
cargo run -- rig --self-test # 6 check groups; should pass
```
## First real run
```bash
cargo run -- all --project-dir out --video day1.mp4 \
--mask-model yolo11x-seg.onnx --mask-device rocm --mask-imgsz 1280
```
- `--mask-device rocm` runs YOLO on an AMD GPU (see [GPU](GPU.md)); use `cpu`
if you don't have the ROCm/MIGraphX runtime.
- The first GPU masks run compiles the model (slow); later runs are fast (see
the cache note in [GPU](GPU.md)).
+14
View File
@@ -0,0 +1,14 @@
# pipeline_360 — Wiki
Everything beyond the short [README](../README.md), in one place.
## Contents
- **[Getting Started](Getting-Started.md)** — prerequisites, build, first run
- **[Usage](Usage.md)** — CLI reference and the native GUI
- **[Pipeline](Pipeline.md)** — stages, output layout, naming conventions
- **[GPU](GPU.md)** — WebGPU adapter selection, YOLO-on-GPU (MIGraphX), cache,
environment variables, troubleshooting
- **[Architecture](Architecture.md)** — source layout and how a run flows
> The original rewrite plan is still in [`../PLAN.md`](../PLAN.md).
+47
View File
@@ -0,0 +1,47 @@
# Pipeline
The rewrite follows the original Python design but reorders the stages so
**cubemap runs before masks**: masks run on the cube-face images directly,
which avoids a wasteful equirect→cubemap→equirect round-trip.
## Output layout
```
out/
01_frames/ selected sharpest equirect frames (multi-video → <stem> prefix)
02_masks/ (legacy; no longer written by `all`)
03_cubemap/
images/ {stem}_{FACE}.{ext} (equirect → 6 cube faces)
masks/ {stem}_{FACE}.{ext}.png (per-face masks)
04_colmap/
images/ {FACE}/{stem}.{ext} (nested COLMAP tree)
masks/ {FACE}/{stem}.{ext}.png
rig_config.json rig cameras / intrinsics
```
`FACE` ∈ `{F, R, B, L, U, D}`. The nested `{FACE}/…` tree is produced via
symlink / hardlink / copy (`--rig-link-mode`).
## Stage details
1. **Extract** — decodes the video(s) with ffmpeg, computes per-frame GPU
sharpness (variance-of-Laplacian), exposure, and contrast, optionally
filters out blurry / badly exposed / flat frames, checks inter-frame motion
(parallax), then selects the sharpest frame per chunk (or spreads them out).
Writes `_sharpness_metadata_<stem>.csv` so the analysis is reused on
subsequent runs.
2. **Cubemap** — projects each equirect frame to six square cube faces on the
GPU (a faithful port of `py360convert.e2c`), batched for throughput.
3. **Masks** — runs YOLO11x-seg on each cube face, synthesizes a binary mask
(COLMAP convention: **black = excluded**), and writes it next to the face
image. See [GPU](GPU.md).
4. **Rig** — scans the cubemap images + masks, writes `rig_config.json`
(PINHOLE intrinsics + per-face `cam_from_rig_rotation`), and links the
nested COLMAP tree.
## Naming / mask convention
- Masks are named `<image_filename>.png` (hence `frame_00001_F.png.png`).
- In a mask, **black pixels are excluded** (movable objects); white pixels are
kept (the static scene used for reconstruction). Masks mirror the image tree
so COLMAP attaches them to the right cameras.
+62
View File
@@ -0,0 +1,62 @@
# Usage
One binary, `pipeline_360`. Invoked with a **subcommand** it runs the CLI; with
**no subcommand** (or `--gui`) it launches the native GUI (requires the `gui`
feature).
## `all` — full pipeline
Runs extract → cubemap → masks → rig on one GPU device:
```bash
pipeline_360 all --project-dir out --video v1.mp4 [v2.mp4 …] \
--mask-model yolo11x-seg.onnx --mask-device rocm --mask-imgsz 1280 \
--extract-chunk-size 30 --cubemap-face-size 1600 --rig-ref-face F
```
- Multi-video: each source is prefixed with its filename stem in `01_frames`.
- `--start-stage extract|cubemap|masks|rig` — resume from a later stage.
- `--skip-masks` — skip the masks stage entirely.
## Individual stages
```bash
# Extract sharpest frames (GPU sharpness + quality filtering)
pipeline_360 extract --video v1.mp4 [v2.mp4 …] --output-dir out/01_frames \
--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
# Cubemap projection (equirect → 6 cube faces, GPU)
pipeline_360 cubemap --frames-dir 01_frames --output-frames-dir 03_cubemap/images \
--face-size 1600 --skip-faces "U,D" --resume --jpeg-quality 95
# Masks on the cube-face images (GPU via --mask-device rocm)
pipeline_360 masks --input-dir 03_cubemap/images --output-dir 03_cubemap/masks \
--mask-model yolo11x-seg.onnx --mask-device rocm --mask-imgsz 1280 \
--mask-classes person,car,bus,truck,bicycle,motorcycle \
--mask-conf 0.35 --mask-dilate 6 --mask-skip-faces "" --mask-empty-if-none
# Rig config + nested COLMAP tree
pipeline_360 rig --images-dir 03_cubemap/images --masks-dir 03_cubemap/masks \
--output rig_config.json --link-tree 04_colmap/images \
--masks-link-tree 04_colmap/masks --face-size 1600 --link-mode symlink
```
## GUI
The native egui+wgpu GUI is in the same binary:
```bash
cargo build --features gui
./target/debug/pipeline_360 # no subcommand → GUI
./target/debug/pipeline_360 --gui # explicit
```
The GUI shells out to the CLI subprocess (the same binary) with
`--gui-log-output` and streams progress into a live log panel. It exposes every
stage's options as forms.
> `--mask-projection` and `--mask-face-size` are accepted for back-compat but
> are no-ops: masks always run on cube-face images.