Ben_Kosytorz/3sixty2cubemap
SHA256
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Compare commits

base: Ben_Kosytorz/3sixty2cubemap:e5b7863320e6027daa54935542aa64a976ee400098c75dff4765f9b183296779
Branches Tags
Ben_Kosytorz/3sixty2cubemap:main Ben_Kosytorz/3sixty2cubemap:rewrite Ben_Kosytorz/3sixty2cubemap:python_old
..
compare: Ben_Kosytorz/3sixty2cubemap:rewrite
Branches Tags
Ben_Kosytorz/3sixty2cubemap:rewrite Ben_Kosytorz/3sixty2cubemap:main Ben_Kosytorz/3sixty2cubemap:python_old
3 Commits
e5b7863320 .. rewrite
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
12 changed files with 1259 additions and 455 deletions
Expand all files Collapse all files
Cargo.toml
+2 -1
View File
@@ -3,7 +3,8 @@ name = "pipeline-360"
version = "0.1.0"
edition = "2021"
description = "360° equirect -> COLMAP cubemap rig pipeline (Rust/WebGPU rewrite)"
license = "MIT"
# The full GPL-3.0 text lives in ./LICENSE.
license = "GPL-3.0-or-later"
[lib]
name = "pipeline_360"
README.md
+59 -207
View File
@@ -1,226 +1,78 @@
# pipeline_360 — Rust/WebGPU rewrite
# pipeline_360
An in-progress rewrite of the Python `pipeline_360.py` equirect→COLMAP-cubemap
rig pipeline in **Rust**, with all heavy image math on the **GPU through
WebGPU** (`wgpu`) for cross-vendor compatibility (Vulkan / Metal / DX12, and
WebGPU/wasm).
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.
> **Status: complete — single binary, masks default, native GUI, GPU masks.**
> One `pipeline_360` binary runs the full GPU pipeline and, with
> `--features gui`, also launches a native egui+wgpu GUI when invoked with no
> subcommand. The YOLO masks stage is on by default. Validated end-to-end on
> real 8K equirect video. `cargo test` passes; 0 clippy warnings in pipeline
> code. YOLO11x-seg via ONNX Runtime, and masks now run on the **GPU** (AMD
> MIGraphX/ROCm) and on the **cube-face images** directly (no equirect
> round-trip). All GPU stages are **batched** for throughput.
> See [PLAN.md](PLAN.md) for the full design and milestone list.
This is the **Rust + WebGPU rewrite** of the original Python pipeline (now in
[`../python`](../python), kept as legacy reference).
## Current state
> ⚡ **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.
| Piece | Status |
| --- | --- |
| `math.rs` (face rotations, quaternions, intrinsics) | ✅ pure std, verified |
| `stages/rig.rs` (rig_config.json + nested tree + linking) | ✅ pure std, verified |
| `selftest.rs` (mirrors Python `rig --self-test`, 6 groups) | ✅ passes (`cargo test`) |
| CLI / `main.rs` (derive-based, mirrors Python commands) | ✅ `extract` + `cubemap` + `all` + `rig` work |
| `video/mod.rs` (ffmpeg/ffprobe CLI decode → rawvideo RGBA) | ✅ verified end-to-end |
| `gpu/mod.rs` (adapter prefers **discrete** GPU; batched readback) | ✅ picks the dGPU (was defaulting to LowPower/iGPU) |
| `gpu/shaders/equirect_to_cubemap.wgsl` | ✅ projection verified vs `py360convert.e2c` |
| `gpu/shaders/sharpness.wgsl` (workgroup-reduced metrics) | ✅ verified end-to-end |
| `gpu/shaders/c2e.wgsl` (cubemap→equirect nearest) | ✅ round-trip validated (legacy mask reprojection) |
| `stages/extract.rs` (GPU sharpness + quality + selection, **batched**) | ✅ superset of Python; verified |
| `stages/cubemap.rs` (GPU equirect→6 faces, skip/resume, **batched**) | ✅ verified vs `py360convert.e2c` |
| `stages/masks.rs` (YOLO11x-seg on cube faces; **MIGraphX GPU**; cache) | ✅ IoU 0.96 vs ultralytics |
| `gui.rs` (native **egui+wgpu** GUI; in the same binary via `--features gui`) | ✅ builds & runs |
## What it does
## Build
There is **one binary**, `pipeline_360`. The YOLO masks stage is **on by
default** (via ONNX Runtime). Use `--no-default-features` for a lighter build
without masks.
```bash
cargo build # default: full pipeline incl. YOLO masks
./target/debug/pipeline_360 all --project-dir <proj> --video <v> --skip-masks
cargo test # runs the rig self-test (6 groups)
cargo run -- rig --self-test
# Lighter build without ONNX Runtime (no masks stage):
cargo build --no-default-features
# Run the GPU extract stage (needs ffmpeg/ffprobe on PATH + a GPU):
cargo run -- extract --video <video> --output-dir <out> \
[--chunk-size 30] [--scale-width 1920] [--output-format png]
[--min-sharpness 0] [--min-exposure 0 --max-exposure 255]
[--min-contrast 0] [--min-motion 0 --max-motion 1e12]
[--selection chunk|spread] [--target-count 0] [--min-gap 0]
[--bad-frame-action keep|skip|warn] [--motion-mode frame_diff|off]
# Run the GPU cubemap stage (needs a Vulkan/Metal/DX12 device):
cargo run -- cubemap --frames-dir <equirect_dir> --output-frames-dir <out> \
[--face-size 1600] [--skip-faces "U,D"] [--resume]
[--masks-dir <dir> --output-masks-dir <mask_out>] [--jpeg-quality 95]
# Native GUI (egui + wgpu; cross-vendor/cross-platform) — same binary, with the `gui` feature:
cargo build --features gui
./target/debug/pipeline_360 # no subcommand -> launches the GUI
./target/debug/pipeline_360 --gui # explicit
./target/debug/pipeline_360 all --project-dir <proj> --video <v> # subcommand -> CLI
# Full pipeline (extract → cubemap → masks → rig), as the GUI invokes it.
# Note the order: cubemap runs *before* masks so masks run on the cube-face
# images (no wasteful equirect round-trip).
cargo run -- all --project-dir <proj> --video <v1> [<v2> …] \
--extract-chunk-size 30 --cubemap-face-size 1600 --rig-ref-face F \
--mask-model yolo11x-seg.onnx --mask-device rocm --mask-imgsz 1280
# YOLO masks stage (runs on the cube-face images the cubemap stage produced):
# Export once: yolo export model=yolo11x-seg.pt format=onnx imgsz=1280 dynamic=True
cargo run -- masks --input-dir <03_cubemap/images> --output-dir <03_cubemap/masks> \
[--mask-classes person,car,...] [--mask-model yolo11x-seg.onnx] \
[--mask-device cpu|rocm] [--mask-imgsz 1280] [--mask-conf 0.35] \
[--mask-dilate 6] [--mask-skip-faces ""] [--mask-empty-if-none] \
[--mask-ep-cache <dir>] # MIGraphX compiled-model cache dir (see below)
# `--mask-projection` / `--mask-face-size` are accepted for CLI/GUI backwards
# compatibility but are now no-ops: masks always run on cube-face images.
```
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)
```
> Decode uses the **`ffmpeg` / `ffprobe` CLI** (piped `rawvideo`), so no native
> FFmpeg bindings or `video` feature are required — just have `ffmpeg` on PATH.
> Fetching crates and building `wgpu` requires network access. In this
> environment that was granted, so `cargo build` / `cargo test` pass, and the
> extract + cubemap stages were verified end-to-end.
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.
## What the extract stage adds over the original Python
The original only picked the sharpest frame (variance-of-Laplacian) per fixed
chunk. The Rust extract stage keeps that as the default and adds genuine
quality control, useful for 360° SfM:
| Flag | Purpose |
| --- | --- |
| `--min-sharpness` | Drop blurry frames (variance-of-Laplacian floor) |
| `--min-exposure` / `--max-exposure` | Drop over/under-exposed frames (mean luma 0..255) |
| `--min-contrast` | Drop flat/featureless frames (luma std floor) |
| `--min-motion` | **Parallax check** — require inter-frame motion (baseline) between kept frames |
| `--max-motion` | Reject jittery/too-fast frames (likely motion blur) |
| `--selection spread` | Pick well-distributed frames instead of fixed chunks |
| `--target-count N` | Target exactly N frames (spread mode) |
| `--min-gap` | Enforce a minimum frame spacing between selected frames |
| `--bad-frame-action skip|keep|warn` | What to do when a chunk has no surviving frame |
| `--motion-mode off|frame_diff` | Toggle the inter-frame motion proxy (cheap thumbnail diff) |
Metrics are written to `_sharpness_metadata[_<stem>].csv` as
`frame,sharpness,exposure,contrast,motion` and reused on subsequent runs unless
`--extract-regenerate-metadata` is given.
### Batching
`--extract-batch` (default 8) and `--cubemap-batch` (default 4) submit a group
of frames' GPU work in one command encoder and read it all back with a single
`device.poll` — amortising per-frame readback latency and letting the GPU
pipeline across frames. Output is identical to the non-batched path (validated
for both stages).
### Masks stage (GPU, on cube faces)
YOLO11x-seg segmentation via ONNX Runtime (`--features masks`). The masks
stage runs **after** the cubemap stage and operates directly on the cube-face
images it produced (`03_cubemap/images/{stem}_{FACE}.ext`). Rationale:
- Masking equirect frames directly is useless (poles distort objects).
- The old design round-tripped `e2c → YOLO → c2e → equirect mask`, then the
cubemap stage did `e2c` **again** — two projections + a stitch-back, wasted.
- Now the cubemap stage projects once; masks run on the square, undistorted cube
faces and stay as cube-face masks (`03_cubemap/masks/{stem}_{FACE}.ext.png`).
- Because the inputs are square, the only resize is a uniform scale to `imgsz`
with **no letterbox padding** — so there is no crop/offset to mis-scale on the
way back. (The previous cubemap path used `r=imgsz/face_size` but letterboxed
with `r=1.0` + a centre pad, producing **shifted masks**; that bug is gone by
construction. Validated at IoU 0.96 vs ultralytics on a real cube face.)
Masks use the COLMAP convention: **black = excluded**. Movable objects
(person, car, …) are masked out (black); the static scene is kept (white).
**GPU execution.** `ort` is built with `load-dynamic`, so it `dlopen`s the
*system* `libonnxruntime` at runtime (override with `ORT_DYLIB_PATH`) instead
of downloading a CPU-only build. On AMD/ROCm systems that lib is built with the
**MIGraphX** execution provider, so `--mask-device rocm` runs YOLO on the
Radeon GPU. Falls back to CPU automatically if the GPU EP can't be registered.
## Build & run
```bash
ORT_DYLIB_PATH=/usr/lib/libonnxruntime.so LD_LIBRARY_PATH=/opt/rocm/lib \
cargo run -- masks --input-dir <cube_faces> --output-dir <masks> \
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
```
**Compiled-model cache (auto).** MIGraphX compiles the graph to GPU kernels
on the first run (slow — minutes for a large model on a new arch) and caches the
compiled `.mxr`. The pipeline **auto-sets `ORT_MIGraphX_CACHE_PATH`** to
`~/.cache/pipeline_360/migraphx` (override with `--mask-ep-cache <dir>`), so the
first run compiles+saves and every later run loads in seconds. Without a cache
dir MIGraphX would crash at inference on this build, so the cache is set up for
you automatically.
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`).
```bash
# first run: compiles + saves (slow)
... masks --mask-device rocm ...
# later runs: loads the cached model (fast)
... masks --mask-device rocm ...
# or use a custom cache dir:
... masks --mask-device rocm --mask-ep-cache /some/dir ...
## 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)
```
> The masks stage needs an **ONNX** export of the model (the `.pt` is not
> loadable in Rust): `yolo export model=yolo11x-seg.pt format=onnx imgsz=1280 dynamic=True`.
## License
## Optional offline verification (no cargo)
[GPL-3.0-or-later](LICENSE) — see [`LICENSE`](LICENSE).
The stage-4 core is pure `std`, so it can also be compiled and run with plain
`rustc` as a fallback when crates.io is unreachable:
## More
```bash
rustc --edition 2021 --crate-name pipeline360 offline_check/main.rs \
-o /tmp/pipeline360_selftest && /tmp/pipeline360_selftest
# -> self-test: 6 check groups passed.
```
`offline_check/main.rs` pulls the real source files via `#[path]` and is not
part of the cargo build.
## Native GUI
`gui.rs` is a native **egui** app rendered through **wgpu** (via eframe) —
the same WebGPU layer as the pipeline, so it is cross-vendor/cross-platform
(Vulkan/Metal/DX12). It is part of the **same** `pipeline_360` binary, gated
behind the `gui` feature — run the binary with no subcommand (or `--gui`) and it
launches the GUI; give it a subcommand (`all`, `extract`, …) and it's the CLI:
```bash
cargo build --features gui
./target/debug/pipeline_360 # no subcommand -> GUI
./target/debug/pipeline_360 all --project-dir <proj> --video <v> # subcommand -> CLI
```
The GUI does **not** call the GPU pipeline library in-process. Instead it
shells out to the `pipeline_360` CLI (itself, found beside the GUI exe, else on
PATH) with `--gui-log-output`, streaming the `[progress] ` / log line protocol
on a background thread into a live log panel — the same protocol the Python GUI
uses. This keeps the heavy GPU work in the (optionally release-built) CLI and
keeps the UI responsive, and means the GUI works anywhere the binary does.
The UI exposes every stage's options as forms with the CLI defaults: project
dir + multi-video pickers, start stage, skip-masks, and collapsible sections for
extract (incl. all the quality/selection/batch knobs), masks, cubemap, and rig,
plus Run/Stop and a live log with progress.
> Building `--features gui` pulls a second wgpu major version (egui 0.36 uses
> wgpu 30; the pipeline core uses wgpu 25). Both coexist fine because the GUI
> process only renders with eframe's wgpu while the CLI subprocess does the
> compute with its own. For a smaller binary build the GUI in release mode.
## Layout
See `PLAN.md` for the detailed design (architecture, stages, CLI parity,
milestones, risks, verification strategy).
- **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.
src/gpu/mod.rs
+2
View File
@@ -463,6 +463,8 @@ pub mod shaders {
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");
}
src/gpu/shaders/thumbnail.wgsl
+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;
}
src/gui.rs
+412 -122
View File
@@ -68,16 +68,18 @@ struct RunHandle {
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 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 })
Ok(RunHandle {
child,
rx,
started: Instant::now(),
code: None,
})
}
fn poll(&mut self, log: &mut VecDeque<String>) {
@@ -110,7 +112,11 @@ impl RunHandle {
}
}
fn spawn_line_reader<R: std::io::Read + Send + 'static>(mut r: R, tx: Sender<Line>, is_stderr: bool) {
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();
@@ -167,10 +173,22 @@ struct ExtractForm {
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,
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,
}
}
}
@@ -183,35 +201,65 @@ struct MasksForm {
conf: f64,
imgsz: u32,
dilate: u32,
projection: String,
face_size: u32,
skip_faces: String,
tile_batch: u32,
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, projection: "cubemap".into(), face_size: 1024,
skip_faces: String::new(), tile_batch: 8, empty_if_none: false, iou: 0.45,
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 }
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 } }
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 }
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() } }
fn default() -> Self {
Self {
ref_face: "F".into(),
intrinsics_convention: "linspace".into(),
link_mode: "symlink".into(),
}
}
}
// ---------------------------------------------------------------------------
@@ -239,11 +287,23 @@ 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(),
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(),
}
}
}
@@ -269,55 +329,88 @@ impl App {
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(),
"--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-projection", &m.projection,
"--mask-face-size", &m.face_size.to_string(),
"--mask-skip-faces", &m.skip_faces,
"--mask-tile-batch", &m.tile_batch.to_string(),
"--mask-iou", &fmt_f(m.iou),
"--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(),
"--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,
"--rig-ref-face",
&r.ref_face,
"--rig-intrinsics-convention",
&r.intrinsics_convention,
"--rig-link-mode",
&r.link_mode,
]);
Some(cmd)
}
@@ -334,7 +427,9 @@ impl App {
self.status = "Error: choose a project dir and at least one video.".into();
return;
}
let Some(cmd) = self.build_command() else { return };
let Some(cmd) = self.build_command() else {
return;
};
self.log.clear();
self.last_progress.clear();
self.status = "Running…".into();
@@ -355,20 +450,33 @@ impl App {
fn fmt_f(v: impl ToString) -> String {
let s = v.to_string();
if s.eq_ignore_ascii_case("inf") { "1e12".into() } else { s }
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 args: Vec<String> = cmd
.get_args()
.map(|a| {
let a = a.to_string_lossy().to_string();
if a.contains(' ') { format!("\"{a}\"") } else { a }
}).collect();
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| {
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);
}
@@ -382,7 +490,11 @@ impl eframe::App for App {
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.") };
self.status = if code == 0 {
"Pipeline complete.".into()
} else {
format!("Pipeline exited with code {code}. See log.")
};
}
}
}
@@ -396,18 +508,28 @@ impl eframe::App for App {
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() };
} 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() {
if ui
.add_enabled(running, egui::Button::new("⏹ Stop"))
.clicked()
{
self.cancel();
}
if ui.add_enabled(!running, egui::Button::new("▶ Run pipeline")).clicked() {
if ui
.add_enabled(!running, egui::Button::new("▶ Run pipeline"))
.clicked()
{
self.start();
}
});
@@ -443,7 +565,11 @@ impl eframe::App for App {
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.label(
egui::RichText::new(format!("CLI: {}", self.cli_binary_str))
.small()
.color(egui::Color32::GRAY),
);
ui.add_space(8.0);
ui.label("Project directory");
@@ -486,7 +612,7 @@ impl eframe::App for App {
egui::ComboBox::from_id_salt("start_stage")
.selected_text(self.start_stage.clone())
.show_ui(ui, |ui| {
for s in ["extract", "masks", "cubemap", "rig"] {
for s in ["extract", "cubemap", "masks", "rig"] {
ui.selectable_value(&mut self.start_stage, s.into(), s);
}
});
@@ -504,75 +630,239 @@ impl eframe::App for App {
.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();
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());
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::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 2 — Masks (YOLO segmentation; needs --features masks)")
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| {
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 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; }
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", "cuda", "rocm"], &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("Projection"); combo(ui, "mask_proj", &["cubemap", "equirect"], &mut m.projection); ui.end_row();
ui.label("Face size"); ui.add(egui::DragValue::new(&mut m.face_size).range(64..=4096)); ui.end_row();
ui.label("Skip faces"); ui.text_edit_singleline(&mut m.skip_faces); ui.end_row();
ui.label("Tile batch"); ui.add(egui::DragValue::new(&mut m.tile_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();
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();
});
ui.checkbox(&mut m.empty_if_none, "Write empty (all-keep) mask when no objects detected");
// 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 3 — Cubemap (equirect -> 6 faces)")
.show(ui, |ui| {
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")
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();
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();
});
});
});
src/stages/extract.rs
+400 -93
View File
@@ -40,7 +40,9 @@ impl Selection {
match s {
"chunk" => Ok(Selection::Chunk),
"spread" => Ok(Selection::Spread),
other => Err(anyhow!("unknown --extract-selection `{other}` (chunk|spread)")),
other => Err(anyhow!(
"unknown --extract-selection `{other}` (chunk|spread)"
)),
}
}
}
@@ -62,7 +64,9 @@ impl BadFrameAction {
"skip" => Ok(BadFrameAction::Skip),
"keep" => Ok(BadFrameAction::Keep),
"warn" => Ok(BadFrameAction::Warn),
other => Err(anyhow!("unknown --extract-bad-frame-action `{other}` (skip|keep|warn)")),
other => Err(anyhow!(
"unknown --extract-bad-frame-action `{other}` (skip|keep|warn)"
)),
}
}
}
@@ -80,7 +84,9 @@ impl MotionMode {
match s {
"off" => Ok(MotionMode::Off),
"frame_diff" => Ok(MotionMode::FrameDiff),
other => Err(anyhow!("unknown --extract-motion-mode `{other}` (off|frame_diff)")),
other => Err(anyhow!(
"unknown --extract-motion-mode `{other}` (off|frame_diff)"
)),
}
}
}
@@ -162,6 +168,16 @@ struct AnalyzeParams {
pad1: u32,
}
#[repr(C)]
#[derive(Clone, Copy)]
struct ThumbParams {
w: u32,
h: u32,
tw: u32,
th: u32,
pad0: u32,
}
/// Owns the wgpu pipeline + a pool of reusable resources for *batched*
/// per-frame analysis. `batch` frames are uploaded, dispatched in one encoder
/// and read back with a single `poll(Wait)`, amortising readback latency and
@@ -181,6 +197,14 @@ struct Analyzer<'g> {
width: u32,
height: u32,
batch: usize,
/// GPU motion-thumbnail pass (luma `THUMB_W x THUMB_H` per frame).
thumb_pipeline: wgpu::ComputePipeline,
#[allow(dead_code)]
thumb_bgl: wgpu::BindGroupLayout,
#[allow(dead_code)]
thumb_uniform: wgpu::Buffer,
thumb_bufs: Vec<wgpu::Buffer>,
thumb_bind_groups: Vec<wgpu::BindGroup>,
}
fn pod_bytes<T: Copy>(x: &T) -> &[u8] {
@@ -216,98 +240,259 @@ fn reduce_partial(raw: &[u8]) -> (f64, f64, f64) {
impl<'g> Analyzer<'g> {
fn new(gpu: &'g GpuContext, width: u32, height: u32, batch: usize) -> Result<Self> {
let batch = batch.max(1);
let bgl = gpu.device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
let bgl = gpu
.device
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("analyze-bgl"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0, visibility: wgpu::ShaderStages::COMPUTE,
binding: 0,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform, has_dynamic_offset: false, min_binding_size: None,
}, count: None,
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 1, visibility: wgpu::ShaderStages::COMPUTE,
binding: 1,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: false },
view_dimension: wgpu::TextureViewDimension::D2, multisampled: false,
}, count: None,
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 2, visibility: wgpu::ShaderStages::COMPUTE,
binding: 2,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: false },
has_dynamic_offset: false, min_binding_size: None,
}, count: None,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
],
});
let layout = gpu.device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
let layout = gpu
.device
.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("analyze-layout"),
bind_group_layouts: &[&bgl],
push_constant_ranges: &[],
});
let pipeline = gpu.compute_pipeline("analyze-pipeline", crate::gpu::shaders::SHARPNESS, &layout);
let pipeline =
gpu.compute_pipeline("analyze-pipeline", crate::gpu::shaders::SHARPNESS, &layout);
let n_wg_x = width.div_ceil(16);
let n_wg_y = height.div_ceil(16);
let n_wg = (n_wg_x as u64) * (n_wg_y as u64);
let partial_size = n_wg * 5 * 4;
let uniform = gpu.create_uniform_buffer(
"analyze-params",
pod_bytes(&AnalyzeParams { width, height, pad0: 0, pad1: 0 }),
pod_bytes(&AnalyzeParams {
width,
height,
pad0: 0,
pad1: 0,
}),
);
let mut src_texs = Vec::with_capacity(batch);
let mut partials = Vec::with_capacity(batch);
let mut bind_groups = Vec::with_capacity(batch);
for i in 0..batch {
let tex = gpu.create_texture_2d(&format!("analyze-src-{i}"), width, height);
let partial = gpu.create_storage_buffer(&format!("analyze-partial-{i}"), partial_size, None);
let partial =
gpu.create_storage_buffer(&format!("analyze-partial-{i}"), partial_size, None);
let view = tex.create_view(&wgpu::TextureViewDescriptor::default());
let bg = gpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some(&format!("analyze-bg-{i}")),
layout: &bgl,
entries: &[
wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() },
wgpu::BindGroupEntry { binding: 1, resource: wgpu::BindingResource::TextureView(&view) },
wgpu::BindGroupEntry { binding: 2, resource: partial.as_entire_binding() },
wgpu::BindGroupEntry {
binding: 0,
resource: uniform.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: wgpu::BindingResource::TextureView(&view),
},
wgpu::BindGroupEntry {
binding: 2,
resource: partial.as_entire_binding(),
},
],
});
src_texs.push(tex);
partials.push(partial);
bind_groups.push(bg);
}
Ok(Analyzer { gpu, pipeline, bgl, uniform, src_texs, partials, bind_groups, n_wg_x, n_wg_y, width, height, batch })
// Motion-thumbnail pass: one thread per thumbnail pixel writes its luma
// into a per-frame storage buffer (motion/parallax proxy on the GPU).
let thumb_bgl = gpu
.device
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("thumb-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: false },
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 2,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: false },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
],
});
let thumb_layout = gpu
.device
.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("thumb-layout"),
bind_group_layouts: &[&thumb_bgl],
push_constant_ranges: &[],
});
let thumb_pipeline = gpu.compute_pipeline(
"thumb-pipeline",
crate::gpu::shaders::THUMBNAIL,
&thumb_layout,
);
let thumb_uniform = gpu.create_uniform_buffer(
"thumb-params",
pod_bytes(&ThumbParams {
w: width,
h: height,
tw: THUMB_W,
th: THUMB_H,
pad0: 0,
}),
);
let thumb_size = (THUMB_W * THUMB_H * 4) as u64;
let mut thumb_bufs = Vec::with_capacity(batch);
let mut thumb_bind_groups = Vec::with_capacity(batch);
for i in 0..batch {
let tb = gpu.create_storage_buffer(&format!("thumb-{i}"), thumb_size, None);
let view = src_texs[i].create_view(&wgpu::TextureViewDescriptor::default());
let bg = gpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some(&format!("thumb-bg-{i}")),
layout: &thumb_bgl,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: thumb_uniform.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: wgpu::BindingResource::TextureView(&view),
},
wgpu::BindGroupEntry {
binding: 2,
resource: tb.as_entire_binding(),
},
],
});
thumb_bufs.push(tb);
thumb_bind_groups.push(bg);
}
Ok(Analyzer {
gpu,
pipeline,
bgl,
uniform,
src_texs,
partials,
bind_groups,
n_wg_x,
n_wg_y,
width,
height,
batch,
thumb_pipeline,
thumb_bgl,
thumb_uniform,
thumb_bufs,
thumb_bind_groups,
})
}
/// Upload + dispatch a batch of frames (len <= `batch`) in one encoder.
/// Runs both the sharpness pass and the motion-thumbnail pass per frame.
fn submit_batch(&self, frames: &[&[u8]]) -> Result<()> {
let n = frames.len();
debug_assert!(n <= self.batch);
// Upload all frames first (separate queue writes), then encode dispatches.
for (i, rgba) in frames.iter().enumerate() {
self.gpu.upload_rgba(&self.src_texs[i], self.width, self.height, rgba);
self.gpu
.upload_rgba(&self.src_texs[i], self.width, self.height, rgba);
}
let mut encoder = self.gpu.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
let mut encoder = self
.gpu
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
{
let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("analyze-pass"), timestamp_writes: None,
label: Some("analyze-pass"),
timestamp_writes: None,
});
pass.set_pipeline(&self.pipeline);
for i in 0..n {
pass.set_bind_group(0, &self.bind_groups[i], &[]);
pass.dispatch_workgroups(self.n_wg_x, self.n_wg_y, 1);
}
// Motion-thumbnail pass (read-only src, write luma thumbnail).
pass.set_pipeline(&self.thumb_pipeline);
let thumb_wgs = (THUMB_W * THUMB_H).div_ceil(256);
for i in 0..n {
pass.set_bind_group(0, &self.thumb_bind_groups[i], &[]);
pass.dispatch_workgroups(thumb_wgs, 1, 1);
}
}
self.gpu.queue.submit(Some(encoder.finish()));
Ok(())
}
/// Read back a batch of `n` partial buffers with one poll and reduce.
fn collect_batch(&self, n: usize) -> Result<Vec<(f64, f64, f64)>> {
let size = (self.n_wg_x as u64) * (self.n_wg_y as u64) * 5 * 4;
let srcs: Vec<&wgpu::Buffer> = self.partials[..n].iter().collect();
let sizes = vec![size; n];
/// Read back a batch of `n` partial buffers + thumbnails with one poll.
/// Returns `(sharpness/exposure/contrast, thumbnails)` per frame.
fn collect_batch(&self, n: usize) -> Result<(Vec<(f64, f64, f64)>, Vec<Vec<u8>>)> {
let psize = (self.n_wg_x as u64) * (self.n_wg_y as u64) * 5 * 4;
let mut srcs: Vec<&wgpu::Buffer> = self.partials[..n].iter().collect();
let mut sizes: Vec<u64> = vec![psize; n];
let tsize = (THUMB_W * THUMB_H * 4) as u64;
srcs.extend(self.thumb_bufs[..n].iter());
sizes.extend(vec![tsize; n]);
let raws = self.gpu.readback_many(&srcs, &sizes)?;
Ok(raws.iter().map(|r| reduce_partial(r)).collect())
let mut metrics = Vec::with_capacity(n);
let mut thumbs = Vec::with_capacity(n);
for r in &raws[..n] {
metrics.push(reduce_partial(r));
}
for r in &raws[n..] {
thumbs.push(r.clone());
}
Ok((metrics, thumbs))
}
}
@@ -318,24 +503,6 @@ impl<'g> Analyzer<'g> {
const THUMB_W: u32 = 64;
const THUMB_H: u32 = 32;
/// Build a tiny grayscale thumbnail (THUMB_W x THUMB_H) from RGBA.
fn thumbnail(rgba: &[u8], w: u32, h: u32) -> Vec<u8> {
let mut out = vec![0u8; (THUMB_W * THUMB_H) as usize];
for ty in 0..THUMB_H {
for tx in 0..THUMB_W {
// Nearest-sample into the source frame.
let sx = (u64::from(tx) * u64::from(w) / u64::from(THUMB_W)) as usize;
let sy = (u64::from(ty) * u64::from(h) / u64::from(THUMB_H)) as usize;
let i = (sy * w as usize + sx) * 4;
let r = rgba[i] as f32;
let g = rgba[i + 1] as f32;
let b = rgba[i + 2] as f32;
out[ty as usize * THUMB_W as usize + tx as usize] = (0.299 * r + 0.587 * g + 0.114 * b) as u8;
}
}
out
}
/// Mean absolute luminance difference between two thumbnails, in 0..255.
fn frame_motion(a: &[u8], b: &[u8]) -> f64 {
if a.len() != b.len() {
@@ -377,12 +544,30 @@ fn read_metadata(path: &Path) -> Result<Vec<FrameMetrics>> {
}
let mut it = line.split(',');
let frame = it.next().and_then(|s| s.trim().parse::<u64>().ok());
let sharp = it.next().and_then(|s| s.trim().parse::<f64>().ok()).unwrap_or(0.0);
let exposure = it.next().and_then(|s| s.trim().parse::<f64>().ok()).unwrap_or(0.0);
let contrast = it.next().and_then(|s| s.trim().parse::<f64>().ok()).unwrap_or(0.0);
let motion = it.next().and_then(|s| s.trim().parse::<f64>().ok()).unwrap_or(0.0);
let sharp = it
.next()
.and_then(|s| s.trim().parse::<f64>().ok())
.unwrap_or(0.0);
let exposure = it
.next()
.and_then(|s| s.trim().parse::<f64>().ok())
.unwrap_or(0.0);
let contrast = it
.next()
.and_then(|s| s.trim().parse::<f64>().ok())
.unwrap_or(0.0);
let motion = it
.next()
.and_then(|s| s.trim().parse::<f64>().ok())
.unwrap_or(0.0);
if let Some(frame) = frame {
out.push(FrameMetrics { frame, sharpness: sharp, exposure, contrast, motion });
out.push(FrameMetrics {
frame,
sharpness: sharp,
exposure,
contrast,
motion,
});
}
}
Ok(out)
@@ -431,7 +616,8 @@ fn select_frames(cfg: &ExtractConfig, metrics: &[FrameMetrics]) -> Vec<u64> {
let end = (i + chunk).min(metrics.len());
let window = &metrics[i..end];
// Prefer the sharpest surviving frame; fall back per bad-frame-action.
let survivors: Vec<&FrameMetrics> = window.iter().filter(|m| passes(cfg, m)).collect();
let survivors: Vec<&FrameMetrics> =
window.iter().filter(|m| passes(cfg, m)).collect();
let pick = survivors
.iter()
.max_by(|a, b| a.sharpness.partial_cmp(&b.sharpness).unwrap())
@@ -440,12 +626,10 @@ fn select_frames(cfg: &ExtractConfig, metrics: &[FrameMetrics]) -> Vec<u64> {
Some(m) => Some(m.frame),
None => match cfg.bad_frame_action {
BadFrameAction::Skip => None,
BadFrameAction::Warn | BadFrameAction::Keep => {
window
BadFrameAction::Warn | BadFrameAction::Keep => window
.iter()
.max_by(|a, b| a.sharpness.partial_cmp(&b.sharpness).unwrap())
.map(|m| m.frame)
}
.map(|m| m.frame),
},
};
if let Some(f) = chosen {
@@ -465,7 +649,11 @@ fn select_frames(cfg: &ExtractConfig, metrics: &[FrameMetrics]) -> Vec<u64> {
let mut survivors: Vec<&FrameMetrics> =
metrics.iter().filter(|m| passes(cfg, m)).collect();
survivors.sort_by(|a, b| b.sharpness.partial_cmp(&a.sharpness).unwrap());
let target = if cfg.target_count == 0 { survivors.len() as u32 } else { cfg.target_count };
let target = if cfg.target_count == 0 {
survivors.len() as u32
} else {
cfg.target_count
};
let mut accepted: Vec<u64> = Vec::new();
for m in survivors {
if accepted.len() as u32 >= target {
@@ -490,7 +678,12 @@ fn select_frames(cfg: &ExtractConfig, metrics: &[FrameMetrics]) -> Vec<u64> {
// ===========================================================================
/// Format a printf-style pattern with one integer (supports `%d`, `%05d`, …).
pub fn format_output_filename(pattern: &str, index: u32, output_format: &str, prefix: &str) -> Result<String> {
pub fn format_output_filename(
pattern: &str,
index: u32,
output_format: &str,
prefix: &str,
) -> Result<String> {
// Find the single `%[flags][width]d` token and substitute.
let bytes = pattern.as_bytes();
let mut out = String::new();
@@ -521,9 +714,21 @@ pub fn format_output_filename(pattern: &str, index: u32, output_format: &str, pr
let composed = format!("{prefix}{out}");
let path = Path::new(&composed);
let target_suffix = format!(".{output_format}");
let final_name = if path.extension().map(|e| e.to_string_lossy().to_lowercase()).as_deref() == Some("jpg")
|| path.extension().map(|e| e.to_string_lossy().to_lowercase()).as_deref() == Some("jpeg")
|| path.extension().map(|e| e.to_string_lossy().to_lowercase()).as_deref() == Some("png")
let final_name = if path
.extension()
.map(|e| e.to_string_lossy().to_lowercase())
.as_deref()
== Some("jpg")
|| path
.extension()
.map(|e| e.to_string_lossy().to_lowercase())
.as_deref()
== Some("jpeg")
|| path
.extension()
.map(|e| e.to_string_lossy().to_lowercase())
.as_deref()
== Some("png")
{
path.with_extension(target_suffix.trim_start_matches('.'))
.to_string_lossy()
@@ -540,7 +745,14 @@ pub fn format_output_filename(pattern: &str, index: u32, output_format: &str, pr
// Frame writing
// ===========================================================================
fn save_frame(rgba: &[u8], w: u32, h: u32, path: &Path, jpeg_quality: u32, output_format: &str) -> Result<()> {
fn save_frame(
rgba: &[u8],
w: u32,
h: u32,
path: &Path,
jpeg_quality: u32,
output_format: &str,
) -> Result<()> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
@@ -552,8 +764,10 @@ fn save_frame(rgba: &[u8], w: u32, h: u32, path: &Path, jpeg_quality: u32, outpu
};
if fmt == image::ImageFormat::Jpeg {
let rgb = image::DynamicImage::ImageRgba8(img).to_rgb8();
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);
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(&image::DynamicImage::ImageRgb8(rgb))
.with_context(|| format!("cannot encode jpeg {}", path.display()))?;
} else {
@@ -587,7 +801,7 @@ fn analyze_video(
) -> Result<Vec<FrameMetrics>> {
let (w, h) = video::scaled_dims(probe, cfg.scale_width);
logging::emit(
&format!("Analyzing sharpness @ {w}x{h} (GPU Laplacian + luma stats)…", ),
&format!("Analyzing sharpness @ {w}x{h} (GPU Laplacian + luma stats)…",),
logger,
);
let batch = cfg.batch.max(1) as usize;
@@ -614,18 +828,32 @@ fn analyze_video(
}
let refs: Vec<&[u8]> = buf.iter().map(|v| v.as_slice()).collect();
analyzer.submit_batch(&refs)?;
let results = analyzer.collect_batch(refs.len())?;
for (i, (sharp, exp, contrast)) in results.into_iter().enumerate() {
let rgba = &buf[i];
let (metrics_b, thumbs_b) = analyzer.collect_batch(refs.len())?;
for (i, (sharp, exp, contrast)) in metrics_b.into_iter().enumerate() {
let motion = if want_motion {
let thumb = thumbnail(rgba, w, h);
let m = prev_thumb.as_ref().map(|p| frame_motion(&thumb, p)).unwrap_or(0.0);
prev_thumb = Some(thumb);
// Thumbnails come back from the GPU as an `array<f32>` of luma
// (0..255). Decode each 4-byte element to an f32, then round to
// u8 for the cheap CPU motion diff.
let thumb_u8: Vec<u8> = thumbs_b[i]
.chunks_exact(4)
.map(|px| f32::from_le_bytes([px[0], px[1], px[2], px[3]]).round() as u8)
.collect();
let m = prev_thumb
.as_ref()
.map(|p| frame_motion(&thumb_u8, p))
.unwrap_or(0.0);
prev_thumb = Some(thumb_u8);
m
} else {
0.0
};
metrics.push(FrameMetrics { frame: frame_no, sharpness: sharp, exposure: exp, contrast, motion });
metrics.push(FrameMetrics {
frame: frame_no,
sharpness: sharp,
exposure: exp,
contrast,
motion,
});
frame_no += 1;
}
if frame_no.is_multiple_of(50) || buf.len() < batch {
@@ -660,7 +888,12 @@ fn extract_frames(
std::fs::create_dir_all(output_dir)?;
let targets: std::collections::BTreeSet<u64> = selected.iter().copied().collect();
logging::emit(
&format!("Extracting {} frames at {}x{}…", targets.len(), probe.width, probe.height),
&format!(
"Extracting {} frames at {}x{}…",
targets.len(),
probe.width,
probe.height
),
logger,
);
let mut pipe = FramePipe::open(path, 0, probe)?;
@@ -670,8 +903,20 @@ fn extract_frames(
while let Some(rgba) = pipe.next_frame()? {
should_cancel.check().map_err(|_| anyhow!("Cancelled."))?;
if targets.contains(&frame_no) {
let name = format_output_filename(&cfg.output_pattern, output_index, &cfg.output_format, prefix)?;
save_frame(&rgba, pipe.width, pipe.height, &output_dir.join(name), cfg.jpeg_quality, &cfg.output_format)?;
let name = format_output_filename(
&cfg.output_pattern,
output_index,
&cfg.output_format,
prefix,
)?;
save_frame(
&rgba,
pipe.width,
pipe.height,
&output_dir.join(name),
cfg.jpeg_quality,
&cfg.output_format,
)?;
output_index += 1;
saved += 1;
logging::emit_progress(&format!("extracted {saved}/{}", targets.len()), logger);
@@ -683,7 +928,11 @@ fn extract_frames(
}
pipe.close();
if saved != targets.len() {
bail!("Some selected frames could not be extracted. Expected {}, saved {}.", targets.len(), saved);
bail!(
"Some selected frames could not be extracted. Expected {}, saved {}.",
targets.len(),
saved
);
}
Ok(saved)
}
@@ -693,10 +942,17 @@ fn per_video_prefixes(videos: &[PathBuf]) -> Vec<String> {
let mut used: std::collections::HashMap<String, u32> = std::collections::HashMap::new();
let mut out = Vec::with_capacity(videos.len());
for v in videos {
let base = v.file_stem().map(|s| s.to_string_lossy().to_string()).unwrap_or_default();
let base = v
.file_stem()
.map(|s| s.to_string_lossy().to_string())
.unwrap_or_default();
let n = used.entry(base.clone()).or_insert(0);
*n += 1;
out.push(if *n == 1 { base } else { format!("{base}_{}", *n) });
out.push(if *n == 1 {
base
} else {
format!("{base}_{}", *n)
});
}
out
}
@@ -733,7 +989,15 @@ pub fn run_extract_stage(
for (idx, (video, prefix)) in videos.iter().zip(prefixes.iter()).enumerate() {
should_cancel.check().map_err(|_| anyhow!("Cancelled."))?;
if multi {
logging::emit(&format!("--- Source {}/{}: {} ---", idx + 1, videos.len(), video.display()), logger);
logging::emit(
&format!(
"--- Source {}/{}: {} ---",
idx + 1,
videos.len(),
video.display()
),
logger,
);
}
let tag = if multi { prefix.as_str() } else { "" };
let metadata_path = output_dir.join(if multi {
@@ -744,18 +1008,36 @@ pub fn run_extract_stage(
let probe = video::probe(video)?;
logging::emit(
&format!("Video: {}x{} {} fps, {} frames declared",
probe.width, probe.height,
probe.fps_f64().map(|f| format!("{f:.3}")).unwrap_or_else(|| "?".into()),
probe.nb_frames.map(|n| n.to_string()).unwrap_or_else(|| "?".into())),
&format!(
"Video: {}x{} {} fps, {} frames declared",
probe.width,
probe.height,
probe
.fps_f64()
.map(|f| format!("{f:.3}"))
.unwrap_or_else(|| "?".into()),
probe
.nb_frames
.map(|n| n.to_string())
.unwrap_or_else(|| "?".into())
),
logger,
);
let metrics = if metadata_path.exists() && cfg.reuse_metadata {
logging::emit(&format!("Using existing sharpness metadata: {}", metadata_path.display()), logger);
logging::emit(
&format!(
"Using existing sharpness metadata: {}",
metadata_path.display()
),
logger,
);
read_metadata(&metadata_path)?
} else {
logging::emit(&format!("Generating sharpness metadata: {}", metadata_path.display()), logger);
logging::emit(
&format!("Generating sharpness metadata: {}", metadata_path.display()),
logger,
);
let m = analyze_video(gpu, video, &probe, cfg, should_cancel, logger)?;
write_metadata(&metadata_path, &m)?;
m
@@ -766,21 +1048,46 @@ pub fn run_extract_stage(
let n_rejected = metrics.len() - n_survivors;
if n_rejected > 0 {
logging::emit(
&format!("Quality filter: {n_rejected}/{} frames rejected (thresholds).", metrics.len()),
&format!(
"Quality filter: {n_rejected}/{} frames rejected (thresholds).",
metrics.len()
),
logger,
);
}
logging::emit(
&format!("Selected {}/{} frames (mode={:?}, chunk={}, gap={}).",
selected.len(), metrics.len(), cfg.selection, cfg.chunk_size, cfg.min_gap),
&format!(
"Selected {}/{} frames (mode={:?}, chunk={}, gap={}).",
selected.len(),
metrics.len(),
cfg.selection,
cfg.chunk_size,
cfg.min_gap
),
logger,
);
let saved = extract_frames(video, &probe, &selected, output_dir, cfg, tag, should_cancel, logger)?;
outcomes.push(ExtractOutcome { video: video.to_path_buf(), prefix: prefix.clone(), frames_written: saved });
let saved = extract_frames(
video,
&probe,
&selected,
output_dir,
cfg,
tag,
should_cancel,
logger,
)?;
outcomes.push(ExtractOutcome {
video: video.to_path_buf(),
prefix: prefix.clone(),
frames_written: saved,
});
}
logging::emit("Done: sharp frames extracted.", logger);
logging::emit(&format!("Output directory: {}", output_dir.display()), logger);
logging::emit(
&format!("Output directory: {}", output_dir.display()),
logger,
);
Ok(outcomes)
}
wiki/Architecture.md
+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.
wiki/GPU.md
+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.
wiki/Getting-Started.md
+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)).
wiki/Home.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).
wiki/Pipeline.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.
wiki/Usage.md
+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.