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

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

Browse Source 
systems
This commit is contained in:
Ben_Kosytorz
2026-08-22 16:40:11 +02:00
commit a44539b6db
24 changed files with 7078 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/error.rs
+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())
}
}
src/gpu/mod.rs
+468
View File
@@ -0,0 +1,468 @@
//! 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");
/// Cubemap -> equirect (nearest, for mask reprojection).
pub const C2E: &str = include_str!("shaders/c2e.wgsl");
}
src/gpu/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));
}
src/gpu/shaders/equirect_to_cubemap.wgsl
+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));
}
src/gpu/shaders/sharpness.wgsl
+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];
}
}
src/gui.rs
+600
View File
@@ -0,0 +1,600 @@
//! 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,
projection: String,
face_size: u32,
skip_faces: String,
tile_batch: u32,
empty_if_none: bool,
iou: f64,
}
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,
}
}
}
#[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-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),
]);
if m.empty_if_none {
cmd.arg("--mask-empty-if-none");
}
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", "masks", "cubemap", "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 2 — Masks (YOLO segmentation; needs --features masks)")
.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", "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.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| {
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())),
)
}
src/json.rs
+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));
}
}
src/lib.rs
+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;
src/logging.rs
+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);
}
src/main.rs
+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
}
}
}
src/math.rs
+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)
}
src/project.rs
+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)
}
src/selftest.rs
+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");
}
}
src/stages/cubemap.rs
+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()),
))
}
src/stages/extract.rs
+786
View File
@@ -0,0 +1,786 @@
//! Stage 1 — extract sharpest / best frames per chunk from video.
//!
//! A substantial superset of the Python `run_extract_stage`:
//! * GPU quality metrics (one dispatch/frame, tiny readback) — sharpness
//! (variance-of-Laplacian, the original metric), exposure (mean luma),
//! contrast (std luma).
//! * CPU inter-frame motion proxy (tiny-thumbnail mean-abs-diff) for
//! parallax/baseline checks.
//! * Threshold filtering ("throw away bad frames"): min-sharpness, exposure
//! bounds, min-contrast, min/max-motion.
//! * Two selection modes: `chunk` (original: sharpest per window) and
//! `spread` (target-count, well-distributed, sharpness-maximising).
//! * `--extract-min-gap` enforces a minimum frame spacing between selected
//! frames (parallax baseline), and `--extract-bad-frame-action` decides
//! what happens when a chunk has no surviving frame.
//!
//! Defaults preserve the original Python behaviour (sharpest per chunk, no
//! quality filtering). All new options are opt-in.
use std::path::{Path, PathBuf};
use anyhow::{anyhow, bail, Context, Result};
use crate::error::CancelFlag;
use crate::gpu::GpuContext;
use crate::logging::{self, Logger};
use crate::video::{self, FramePipe, VideoProbe};
/// Selection strategy.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Selection {
/// Sharpest surviving frame per `chunk_size` window (Python default).
Chunk,
/// Up to `target_count` sharp, well-spaced frames across the whole video.
Spread,
}
impl Selection {
pub fn parse(s: &str) -> Result<Self> {
match s {
"chunk" => Ok(Selection::Chunk),
"spread" => Ok(Selection::Spread),
other => Err(anyhow!("unknown --extract-selection `{other}` (chunk|spread)")),
}
}
}
/// What to do when a chunk has no surviving frame after threshold filtering.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BadFrameAction {
/// Omit the chunk entirely (no frame written for it).
Skip,
/// Fall back to the best frame in the chunk regardless of thresholds.
Keep,
/// Like Skip but log a warning.
Warn,
}
impl BadFrameAction {
pub fn parse(s: &str) -> Result<Self> {
match s {
"skip" => Ok(BadFrameAction::Skip),
"keep" => Ok(BadFrameAction::Keep),
"warn" => Ok(BadFrameAction::Warn),
other => Err(anyhow!("unknown --extract-bad-frame-action `{other}` (skip|keep|warn)")),
}
}
}
/// Motion-proxy computation mode.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MotionMode {
Off,
/// Mean absolute luminance difference between consecutive tiny thumbnails.
FrameDiff,
}
impl MotionMode {
pub fn parse(s: &str) -> Result<Self> {
match s {
"off" => Ok(MotionMode::Off),
"frame_diff" => Ok(MotionMode::FrameDiff),
other => Err(anyhow!("unknown --extract-motion-mode `{other}` (off|frame_diff)")),
}
}
}
/// Full extract-stage configuration. Defaults reproduce the original pipeline.
#[derive(Debug, Clone)]
pub struct ExtractConfig {
pub chunk_size: u32,
pub scale_width: u32,
pub workers: u32, // accepted for CLI parity; GPU is the parallelism
pub output_pattern: String,
pub output_format: String, // "png" | "jpg"
pub jpeg_quality: u32,
pub reuse_metadata: bool,
// --- Quality thresholds (all default to "no filtering"). ---
pub min_sharpness: f64,
pub min_exposure: f64,
pub max_exposure: f64,
pub min_contrast: f64,
pub min_motion: f64,
pub max_motion: f64,
// --- Selection ---
pub selection: Selection,
pub target_count: u32, // 0 = chunk-based
pub min_gap: u32, // min frame-index spacing between selected frames
pub bad_frame_action: BadFrameAction,
pub motion_mode: MotionMode,
pub batch: u32,
}
impl Default for ExtractConfig {
fn default() -> Self {
ExtractConfig {
chunk_size: 30,
scale_width: 1920,
workers: 4,
output_pattern: "output_frame_%05d.png".to_string(),
output_format: "png".to_string(),
jpeg_quality: 95,
reuse_metadata: true,
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: Selection::Chunk,
target_count: 0,
min_gap: 0,
bad_frame_action: BadFrameAction::Keep,
motion_mode: MotionMode::FrameDiff,
batch: 8,
}
}
}
/// Per-frame quality metrics produced by the analysis pass.
#[derive(Debug, Clone, Copy)]
pub struct FrameMetrics {
pub frame: u64,
pub sharpness: f64,
pub exposure: f64,
pub contrast: f64,
pub motion: f64, // 0.0 if motion_mode == Off / first frame
}
// ===========================================================================
// GPU analysis
// ===========================================================================
#[repr(C)]
#[derive(Clone, Copy)]
struct AnalyzeParams {
width: u32,
height: u32,
pad0: u32,
pad1: 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
/// letting the GPU pipeline work across frames.
struct Analyzer<'g> {
gpu: &'g GpuContext,
pipeline: wgpu::ComputePipeline,
#[allow(dead_code)]
bgl: wgpu::BindGroupLayout,
#[allow(dead_code)]
uniform: wgpu::Buffer,
src_texs: Vec<wgpu::Texture>,
partials: Vec<wgpu::Buffer>,
bind_groups: Vec<wgpu::BindGroup>,
n_wg_x: u32,
n_wg_y: u32,
width: u32,
height: u32,
batch: usize,
}
fn pod_bytes<T: Copy>(x: &T) -> &[u8] {
unsafe { std::slice::from_raw_parts((x as *const T) as *const u8, std::mem::size_of::<T>()) }
}
fn reduce_partial(raw: &[u8]) -> (f64, f64, f64) {
let floats: &[f32] =
unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, raw.len() / 4) };
let mut sum_g = 0.0f64;
let mut sum_g2 = 0.0f64;
let mut sum_l = 0.0f64;
let mut sum_l2 = 0.0f64;
let mut n = 0.0f64;
for chunk in floats.as_chunks::<5>().0 {
sum_g += chunk[0] as f64;
sum_g2 += chunk[1] as f64;
sum_l += chunk[2] as f64;
sum_l2 += chunk[3] as f64;
n += chunk[4] as f64;
}
if n <= 0.0 {
return (0.0, 0.0, 0.0);
}
let mean_g = sum_g / n;
let var_g = (sum_g2 / n - mean_g * mean_g).max(0.0);
let contrast = var_g.sqrt();
let mean_l = sum_l / n;
let var_l = (sum_l2 / n - mean_l * mean_l).max(0.0);
(var_l, mean_g, contrast)
}
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 {
label: Some("analyze-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 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 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 }),
);
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 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() },
],
});
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 })
}
/// Upload + dispatch a batch of frames (len <= `batch`) in one encoder.
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);
}
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,
});
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);
}
}
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];
let raws = self.gpu.readback_many(&srcs, &sizes)?;
Ok(raws.iter().map(|r| reduce_partial(r)).collect())
}
}
// ===========================================================================
// CPU motion proxy
// ===========================================================================
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() {
return 0.0;
}
let mut acc = 0u64;
for (x, y) in a.iter().zip(b.iter()) {
acc += (*x as i64 - *y as i64).unsigned_abs();
}
acc as f64 / a.len() as f64
}
// ===========================================================================
// Metadata CSV
// ===========================================================================
fn write_metadata(path: &Path, metrics: &[FrameMetrics]) -> Result<()> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let mut s = String::from("frame,sharpness,exposure,contrast,motion\n");
for m in metrics {
s.push_str(&format!(
"{},{:.10},{:.10},{:.10},{:.10}\n",
m.frame, m.sharpness, m.exposure, m.contrast, m.motion
));
}
std::fs::write(path, s).with_context(|| format!("cannot write metadata {}", path.display()))?;
Ok(())
}
fn read_metadata(path: &Path) -> Result<Vec<FrameMetrics>> {
let txt = std::fs::read_to_string(path)
.with_context(|| format!("cannot read metadata {}", path.display()))?;
let mut out = Vec::new();
for (i, line) in txt.lines().enumerate() {
if i == 0 || line.trim().is_empty() {
continue;
}
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);
if let Some(frame) = frame {
out.push(FrameMetrics { frame, sharpness: sharp, exposure, contrast, motion });
}
}
Ok(out)
}
// ===========================================================================
// Selection
// ===========================================================================
/// True if a frame passes the configured quality thresholds.
fn passes(cfg: &ExtractConfig, m: &FrameMetrics) -> bool {
m.sharpness >= cfg.min_sharpness
&& m.exposure >= cfg.min_exposure
&& m.exposure <= cfg.max_exposure
&& m.contrast >= cfg.min_contrast
&& m.motion >= cfg.min_motion
&& m.motion <= cfg.max_motion
}
/// Select the winning frame indices from per-frame metrics.
fn select_frames(cfg: &ExtractConfig, metrics: &[FrameMetrics]) -> Vec<u64> {
if metrics.is_empty() {
return Vec::new();
}
// First enforce min-gap as a post-filter on any selection.
let enforce_gap = |sel: &mut Vec<u64>| {
if cfg.min_gap == 0 {
return;
}
let mut out: Vec<u64> = Vec::new();
for &f in sel.iter() {
if out.last().is_none_or(|&p| f >= p + u64::from(cfg.min_gap)) {
out.push(f);
}
}
*sel = out;
};
match cfg.selection {
Selection::Chunk => {
let chunk = cfg.chunk_size.max(1) as usize;
let mut sel: Vec<u64> = Vec::new();
let mut i = 0;
while i < metrics.len() {
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 pick = survivors
.iter()
.max_by(|a, b| a.sharpness.partial_cmp(&b.sharpness).unwrap())
.map(|m| **m);
let chosen = match pick {
Some(m) => Some(m.frame),
None => match cfg.bad_frame_action {
BadFrameAction::Skip => None,
BadFrameAction::Warn | BadFrameAction::Keep => {
window
.iter()
.max_by(|a, b| a.sharpness.partial_cmp(&b.sharpness).unwrap())
.map(|m| m.frame)
}
},
};
if let Some(f) = chosen {
if cfg.bad_frame_action == BadFrameAction::Warn && pick.is_none() {
// Caller logs warnings via the count; nothing here.
}
sel.push(f);
}
i = end;
}
enforce_gap(&mut sel);
sel
}
Selection::Spread => {
// Greedy: sort survivors by sharpness desc, accept if >= min_gap from
// all accepted; stop at target_count (or all survivors if target==0).
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 mut accepted: Vec<u64> = Vec::new();
for m in survivors {
if accepted.len() as u32 >= target {
break;
}
let ok = accepted
.iter()
.all(|&p| (m.frame as i64 - p as i64).unsigned_abs() >= u64::from(cfg.min_gap));
if ok {
accepted.push(m.frame);
}
}
// Re-order selected frames by their temporal index for clean output.
accepted.sort_unstable();
accepted
}
}
}
// ===========================================================================
// Output filename formatting (printf-style %0Nd)
// ===========================================================================
/// 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> {
// Find the single `%[flags][width]d` token and substitute.
let bytes = pattern.as_bytes();
let mut out = String::new();
let mut i = 0;
let mut substituted = false;
while i < bytes.len() {
if bytes[i] == b'%' {
// parse %[0-9]*d
let mut j = i + 1;
let mut width_spec = String::new();
while j < bytes.len() && (bytes[j].is_ascii_digit()) {
width_spec.push(bytes[j] as char);
j += 1;
}
if j < bytes.len() && bytes[j] == b'd' && !substituted {
let width: usize = width_spec.parse().unwrap_or(0);
let num = format!("{:0width$}", index, width = width);
out.push_str(&num);
i = j + 1;
substituted = true;
continue;
}
}
out.push(bytes[i] as char);
i += 1;
}
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")
{
path.with_extension(target_suffix.trim_start_matches('.'))
.to_string_lossy()
.to_string()
} else if path.extension().is_some() {
bail!("--extract-output-pattern must end with .png/.jpg or omit the extension");
} else {
format!("{prefix}{out}{target_suffix}")
};
Ok(final_name)
}
// ===========================================================================
// Frame writing
// ===========================================================================
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)?;
}
let img = image::RgbaImage::from_raw(w, h, rgba.to_vec())
.ok_or_else(|| anyhow!("bad frame dims {w}x{h}"))?;
let fmt = match output_format {
"jpg" | "jpeg" => image::ImageFormat::Jpeg,
_ => image::ImageFormat::Png,
};
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);
enc.encode_image(&image::DynamicImage::ImageRgb8(rgb))
.with_context(|| format!("cannot encode jpeg {}", path.display()))?;
} else {
image::DynamicImage::ImageRgba8(img)
.save(path)
.with_context(|| format!("cannot write {}", path.display()))?;
}
Ok(())
}
// ===========================================================================
// The stage
// ===========================================================================
/// Per-video result for multi-video bookkeeping.
pub struct ExtractOutcome {
pub video: PathBuf,
pub prefix: String,
pub frames_written: usize,
}
/// Analyze a single video: decode downscaled, compute metrics on the GPU,
/// return the per-frame metrics.
fn analyze_video(
gpu: &GpuContext,
path: &Path,
probe: &VideoProbe,
cfg: &ExtractConfig,
should_cancel: &CancelFlag,
logger: Logger,
) -> 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)…", ),
logger,
);
let batch = cfg.batch.max(1) as usize;
let analyzer = Analyzer::new(gpu, w, h, batch)?;
let mut pipe = FramePipe::open(path, cfg.scale_width, probe)?;
let want_motion = cfg.motion_mode == MotionMode::FrameDiff;
let mut prev_thumb: Option<Vec<u8>> = None;
let mut metrics = Vec::new();
let mut frame_no: u64 = 0;
// Buffered batch: collect up to `batch` decoded frames, submit them all in
// one encoder, read back with one poll, then reduce + compute motion.
let mut buf: Vec<Vec<u8>> = Vec::with_capacity(batch);
loop {
should_cancel.check().map_err(|_| anyhow!("Cancelled."))?;
while buf.len() < batch {
match pipe.next_frame()? {
Some(rgba) => buf.push(rgba),
None => break,
}
}
if buf.is_empty() {
break;
}
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 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);
m
} else {
0.0
};
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 {
logging::emit_progress(&format!("analyzed {frame_no} frames"), logger);
}
buf.clear();
}
pipe.close();
if metrics.is_empty() {
bail!("No frames could be read from the video.");
}
logging::emit(&format!("Analyzed {} frames.", metrics.len()), logger);
Ok(metrics)
}
/// Extract the selected frames at full resolution.
#[allow(clippy::too_many_arguments)]
fn extract_frames(
path: &Path,
probe: &VideoProbe,
selected: &[u64],
output_dir: &Path,
cfg: &ExtractConfig,
prefix: &str,
should_cancel: &CancelFlag,
logger: Logger,
) -> Result<usize> {
if selected.is_empty() {
logging::emit("No frames selected for extraction.", logger);
return Ok(0);
}
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),
logger,
);
let mut pipe = FramePipe::open(path, 0, probe)?;
let mut frame_no: u64 = 0;
let mut output_index: u32 = 1;
let mut saved: usize = 0;
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)?;
output_index += 1;
saved += 1;
logging::emit_progress(&format!("extracted {saved}/{}", targets.len()), logger);
if saved == targets.len() {
break;
}
}
frame_no += 1;
}
pipe.close();
if saved != targets.len() {
bail!("Some selected frames could not be extracted. Expected {}, saved {}.", targets.len(), saved);
}
Ok(saved)
}
/// Collision-safe per-source prefix (mirrors Python).
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 n = used.entry(base.clone()).or_insert(0);
*n += 1;
out.push(if *n == 1 { base } else { format!("{base}_{}", *n) });
}
out
}
/// Run the extract stage over one or more videos.
#[allow(clippy::too_many_arguments)]
pub fn run_extract_stage(
videos: &[PathBuf],
output_dir: &Path,
cfg: &ExtractConfig,
gpu: &GpuContext,
should_cancel: &CancelFlag,
logger: Logger,
) -> Result<Vec<ExtractOutcome>> {
if videos.is_empty() {
bail!("No input video(s) were provided.");
}
for v in videos {
if !v.exists() {
bail!("Input video was not found: {}", v.display());
}
}
if cfg.chunk_size == 0 {
bail!("--extract-chunk-size must be 1 or greater.");
}
if !matches!(cfg.output_format.as_str(), "png" | "jpg" | "jpeg") {
bail!("--extract-output-format must be png or jpg.");
}
std::fs::create_dir_all(output_dir)?;
let multi = videos.len() > 1;
let prefixes = per_video_prefixes(videos);
let mut outcomes = Vec::new();
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);
}
let tag = if multi { prefix.as_str() } else { "" };
let metadata_path = output_dir.join(if multi {
format!("_sharpness_metadata_{tag}.csv")
} else {
"_sharpness_metadata.csv".to_string()
});
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())),
logger,
);
let metrics = if metadata_path.exists() && cfg.reuse_metadata {
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);
let m = analyze_video(gpu, video, &probe, cfg, should_cancel, logger)?;
write_metadata(&metadata_path, &m)?;
m
};
let selected = select_frames(cfg, &metrics);
let n_survivors = metrics.iter().filter(|m| passes(cfg, m)).count();
let n_rejected = metrics.len() - n_survivors;
if n_rejected > 0 {
logging::emit(
&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),
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 });
}
logging::emit("Done: sharp frames extracted.", logger);
logging::emit(&format!("Output directory: {}", output_dir.display()), logger);
Ok(outcomes)
}
src/stages/masks.rs
+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"
))
}
src/stages/mod.rs
+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;
src/stages/rig.rs
+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",
}
}
src/video/mod.rs
+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();
}
}