3 Commits
5 changed files with 252 additions and 884 deletions
+4 -1
View File
@@ -1,3 +1,6 @@
# Warning
**This project is fully vibecoded and behind no line in this code is human soul.**
# pipeline_360 # pipeline_360
Convert equirectangular 360° videos into a **COLMAP-ready cubemap camera-rig Convert equirectangular 360° videos into a **COLMAP-ready cubemap camera-rig
@@ -5,7 +8,7 @@ dataset** — sharpest-frame extraction, movable-object masks, an
equirect→cubemap projection, and a `rig_config.json` + nested file tree. equirect→cubemap projection, and a `rig_config.json` + nested file tree.
This is the **Rust + WebGPU rewrite** of the original Python pipeline (now in This is the **Rust + WebGPU rewrite** of the original Python pipeline (now in
[`../python`](../python), kept as legacy reference). [python](https://git.dieterunterstrich.de/Ben_Kosytorz/3sixty2cubemap/src/branch/python_old), kept as legacy reference).
> ⚡ **Status: fully vibecoded, and the rewrite is done.** Everything described > ⚡ **Status: fully vibecoded, and the rewrite is done.** Everything described
> here is implemented and running on real 8K test footage — GPU sharpness, > here is implemented and running on real 8K test footage — GPU sharpness,
-2
View File
@@ -463,8 +463,6 @@ pub mod shaders {
pub const EQUIRECT_TO_CUBEMAP: &str = include_str!("shaders/equirect_to_cubemap.wgsl"); pub const EQUIRECT_TO_CUBEMAP: &str = include_str!("shaders/equirect_to_cubemap.wgsl");
/// Downscale + grayscale + Laplacian variance (sharpness scoring). /// Downscale + grayscale + Laplacian variance (sharpness scoring).
pub const SHARPNESS: &str = include_str!("shaders/sharpness.wgsl"); pub const SHARPNESS: &str = include_str!("shaders/sharpness.wgsl");
/// Downscale to a tiny luma thumbnail (motion/parallax proxy).
pub const THUMBNAIL: &str = include_str!("shaders/thumbnail.wgsl");
/// Cubemap -> equirect (nearest, for mask reprojection). /// Cubemap -> equirect (nearest, for mask reprojection).
pub const C2E: &str = include_str!("shaders/c2e.wgsl"); pub const C2E: &str = include_str!("shaders/c2e.wgsl");
} }
-36
View File
@@ -1,36 +0,0 @@
// Stage-1 companion kernel: build a tiny grayscale luma thumbnail on the GPU.
//
// One thread per thumbnail pixel. Each thread nearest-samples a source pixel
// and writes its luma (0..255) into `thumb[ty * tw + tx]`. This lets the motion
// (parallax) proxy be computed on the GPU instead of the CPU reading the full
// RGBA frame, so extraction leans on the GPU rather than the CPU.
struct Params {
w: u32,
h: u32,
tw: u32,
th: u32,
pad0: u32,
}
@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var src_tex: texture_2d<f32>;
@group(0) @binding(2) var<storage, read_write> thumb: array<f32>;
fn luma(c: vec4<f32>) -> f32 {
return 0.299 * c.r + 0.587 * c.g + 0.114 * c.b;
}
@compute @workgroup_size(256, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
let i = gid.x;
if (i >= params.tw * params.th) {
return;
}
let tx = i % params.tw;
let ty = i / params.tw;
let sx = i32(tx * params.w / params.tw);
let sy = i32(ty * params.h / params.th);
// Scale to 0..255 to match the CPU thumbnail / motion-diff convention.
thumb[i] = luma(textureLoad(src_tex, vec2<i32>(sx, sy), 0)) * 255.0;
}
+134 -424
View File
@@ -68,18 +68,16 @@ struct RunHandle {
impl RunHandle { impl RunHandle {
fn spawn(mut cmd: std::process::Command) -> std::io::Result<RunHandle> { fn spawn(mut cmd: std::process::Command) -> std::io::Result<RunHandle> {
let mut child = cmd.stdout(Stdio::piped()).stderr(Stdio::piped()).spawn()?; let mut child = cmd
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()?;
let stdout = child.stdout.take().unwrap(); let stdout = child.stdout.take().unwrap();
let stderr = child.stderr.take().unwrap(); let stderr = child.stderr.take().unwrap();
let (tx, rx) = mpsc::channel::<Line>(); let (tx, rx) = mpsc::channel::<Line>();
spawn_line_reader(stdout, tx.clone(), false); spawn_line_reader(stdout, tx.clone(), false);
spawn_line_reader(stderr, tx, true); spawn_line_reader(stderr, tx, true);
Ok(RunHandle { Ok(RunHandle { child, rx, started: Instant::now(), code: None })
child,
rx,
started: Instant::now(),
code: None,
})
} }
fn poll(&mut self, log: &mut VecDeque<String>) { fn poll(&mut self, log: &mut VecDeque<String>) {
@@ -112,11 +110,7 @@ impl RunHandle {
} }
} }
fn spawn_line_reader<R: std::io::Read + Send + 'static>( fn spawn_line_reader<R: std::io::Read + Send + 'static>(mut r: R, tx: Sender<Line>, is_stderr: bool) {
mut r: R,
tx: Sender<Line>,
is_stderr: bool,
) {
std::thread::spawn(move || { std::thread::spawn(move || {
let mut buf = [0u8; 8192]; let mut buf = [0u8; 8192];
let mut leftover = Vec::new(); let mut leftover = Vec::new();
@@ -173,22 +167,10 @@ struct ExtractForm {
impl Default for ExtractForm { impl Default for ExtractForm {
fn default() -> Self { fn default() -> Self {
Self { Self {
chunk_size: 30, chunk_size: 30, scale_width: 1920, output_format: "png".into(), jpeg_quality: 95,
scale_width: 1920, min_sharpness: 0.0, min_exposure: 0.0, max_exposure: 255.0, min_contrast: 0.0,
output_format: "png".into(), min_motion: 0.0, max_motion: 1.0e12, selection: "chunk".into(), target_count: 0,
jpeg_quality: 95, min_gap: 0, bad_frame_action: "keep".into(), motion_mode: "frame_diff".into(), batch: 8,
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,
} }
} }
} }
@@ -201,65 +183,35 @@ struct MasksForm {
conf: f64, conf: f64,
imgsz: u32, imgsz: u32,
dilate: u32, dilate: u32,
projection: String,
face_size: u32,
skip_faces: String, skip_faces: String,
batch: u32, tile_batch: u32,
empty_if_none: bool, empty_if_none: bool,
iou: f64, iou: f64,
// MIGraphX compiled-model cache directory (optional override; defaults to
// ~/.cache/pipeline_360/migraphx when unset).
ep_cache: Option<PathBuf>,
} }
impl Default for MasksForm { impl Default for MasksForm {
fn default() -> Self { fn default() -> Self {
Self { Self {
classes: "person,car,bus,truck,bicycle,motorcycle".into(), classes: "person,car,bus,truck,bicycle,motorcycle".into(),
model: PathBuf::from("yolo11x-seg.onnx"), model: PathBuf::from("yolo11x-seg.onnx"), device: "cpu".into(), conf: 0.35,
device: "cpu".into(), imgsz: 1280, dilate: 6, projection: "cubemap".into(), face_size: 1024,
conf: 0.35, skip_faces: String::new(), tile_batch: 8, empty_if_none: false, iou: 0.45,
imgsz: 1280,
dilate: 6,
skip_faces: String::new(),
batch: 8,
empty_if_none: false,
iou: 0.45,
ep_cache: None,
} }
} }
} }
#[derive(Clone)] #[derive(Clone)]
struct CubemapForm { struct CubemapForm { face_size: u32, skip_faces: String, jpeg_quality: u32, batch: u32 }
face_size: u32,
skip_faces: String,
jpeg_quality: u32,
batch: u32,
}
impl Default for CubemapForm { impl Default for CubemapForm {
fn default() -> Self { fn default() -> Self { Self { face_size: 1600, skip_faces: "U,D".into(), jpeg_quality: 95, batch: 4 } }
Self {
face_size: 1600,
skip_faces: "U,D".into(),
jpeg_quality: 95,
batch: 4,
}
}
} }
#[derive(Clone)] #[derive(Clone)]
struct RigForm { struct RigForm { ref_face: String, intrinsics_convention: String, link_mode: String }
ref_face: String,
intrinsics_convention: String,
link_mode: String,
}
impl Default for RigForm { impl Default for RigForm {
fn default() -> Self { fn default() -> Self { Self { ref_face: "F".into(), intrinsics_convention: "linspace".into(), link_mode: "symlink".into() } }
Self {
ref_face: "F".into(),
intrinsics_convention: "linspace".into(),
link_mode: "symlink".into(),
}
}
} }
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
@@ -287,23 +239,11 @@ impl Default for App {
fn default() -> Self { fn default() -> Self {
let cli = resolve_cli_binary(); let cli = resolve_cli_binary();
Self { Self {
project_dir: PathBuf::new(), project_dir: PathBuf::new(), videos: Vec::new(), start_stage: "extract".into(),
videos: Vec::new(), skip_masks: false, extract: ExtractForm::default(), masks: MasksForm::default(),
start_stage: "extract".into(), cubemap: CubemapForm::default(), rig: RigForm::default(),
skip_masks: false, cli_binary_str: cli.as_ref().map(|p| p.display().to_string()).unwrap_or_else(|| "pipeline_360 (NOT FOUND — build it first)".into()),
extract: ExtractForm::default(), cli_binary: cli, log: VecDeque::new(), run: None, status: "Idle".into(), last_progress: String::new(),
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(),
} }
} }
} }
@@ -329,88 +269,55 @@ impl App {
let e = &self.extract; let e = &self.extract;
cmd.args([ cmd.args([
"--extract-chunk-size", "--extract-chunk-size", &e.chunk_size.to_string(),
&e.chunk_size.to_string(), "--extract-scale-width", &e.scale_width.to_string(),
"--extract-scale-width", "--extract-output-format", &e.output_format,
&e.scale_width.to_string(), "--extract-jpeg-quality", &e.jpeg_quality.to_string(),
"--extract-output-format", "--extract-min-sharpness", &fmt_f(e.min_sharpness),
&e.output_format, "--extract-min-exposure", &fmt_f(e.min_exposure),
"--extract-jpeg-quality", "--extract-max-exposure", &fmt_f(e.max_exposure),
&e.jpeg_quality.to_string(), "--extract-min-contrast", &fmt_f(e.min_contrast),
"--extract-min-sharpness", "--extract-min-motion", &fmt_f(e.min_motion),
&fmt_f(e.min_sharpness), "--extract-max-motion", &fmt_f(e.max_motion),
"--extract-min-exposure", "--extract-selection", &e.selection,
&fmt_f(e.min_exposure), "--extract-target-count", &e.target_count.to_string(),
"--extract-max-exposure", "--extract-min-gap", &e.min_gap.to_string(),
&fmt_f(e.max_exposure), "--extract-bad-frame-action", &e.bad_frame_action,
"--extract-min-contrast", "--extract-motion-mode", &e.motion_mode,
&fmt_f(e.min_contrast), "--extract-batch", &e.batch.to_string(),
"--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; let m = &self.masks;
cmd.args([ cmd.args([
"--mask-classes", "--mask-classes", &m.classes,
&m.classes, "--mask-model", &m.model.display().to_string(),
"--mask-model", "--mask-device", &m.device,
&m.model.display().to_string(), "--mask-conf", &fmt_f(m.conf),
"--mask-device", "--mask-imgsz", &m.imgsz.to_string(),
&m.device, "--mask-dilate", &m.dilate.to_string(),
"--mask-conf", "--mask-projection", &m.projection,
&fmt_f(m.conf), "--mask-face-size", &m.face_size.to_string(),
"--mask-imgsz", "--mask-skip-faces", &m.skip_faces,
&m.imgsz.to_string(), "--mask-tile-batch", &m.tile_batch.to_string(),
"--mask-dilate", "--mask-iou", &fmt_f(m.iou),
&m.dilate.to_string(),
"--mask-skip-faces",
&m.skip_faces,
"--mask-batch",
&m.batch.to_string(),
"--mask-iou",
&fmt_f(m.iou),
]); ]);
if m.empty_if_none { if m.empty_if_none {
cmd.arg("--mask-empty-if-none"); cmd.arg("--mask-empty-if-none");
} }
if let Some(cache) = &m.ep_cache {
cmd.arg("--mask-ep-cache").arg(cache);
}
let c = &self.cubemap; let c = &self.cubemap;
cmd.args([ cmd.args([
"--cubemap-face-size", "--cubemap-face-size", &c.face_size.to_string(),
&c.face_size.to_string(), "--cubemap-skip-faces", &c.skip_faces,
"--cubemap-skip-faces", "--cubemap-jpeg-quality", &c.jpeg_quality.to_string(),
&c.skip_faces, "--cubemap-batch", &c.batch.to_string(),
"--cubemap-jpeg-quality",
&c.jpeg_quality.to_string(),
"--cubemap-batch",
&c.batch.to_string(),
]); ]);
let r = &self.rig; let r = &self.rig;
cmd.args([ cmd.args([
"--rig-ref-face", "--rig-ref-face", &r.ref_face,
&r.ref_face, "--rig-intrinsics-convention", &r.intrinsics_convention,
"--rig-intrinsics-convention", "--rig-link-mode", &r.link_mode,
&r.intrinsics_convention,
"--rig-link-mode",
&r.link_mode,
]); ]);
Some(cmd) Some(cmd)
} }
@@ -427,9 +334,7 @@ impl App {
self.status = "Error: choose a project dir and at least one video.".into(); self.status = "Error: choose a project dir and at least one video.".into();
return; return;
} }
let Some(cmd) = self.build_command() else { let Some(cmd) = self.build_command() else { return };
return;
};
self.log.clear(); self.log.clear();
self.last_progress.clear(); self.last_progress.clear();
self.status = "Running…".into(); self.status = "Running…".into();
@@ -450,37 +355,24 @@ impl App {
fn fmt_f(v: impl ToString) -> String { fn fmt_f(v: impl ToString) -> String {
let s = v.to_string(); let s = v.to_string();
if s.eq_ignore_ascii_case("inf") { if s.eq_ignore_ascii_case("inf") { "1e12".into() } else { s }
"1e12".into()
} else {
s
}
} }
fn cmd_to_string(cmd: &std::process::Command) -> String { fn cmd_to_string(cmd: &std::process::Command) -> String {
let prog = cmd.get_program().to_string_lossy().to_string(); let prog = cmd.get_program().to_string_lossy().to_string();
let args: Vec<String> = cmd let args: Vec<String> = cmd.get_args().map(|a| {
.get_args() let a = a.to_string_lossy().to_string();
.map(|a| { if a.contains(' ') { format!("\"{a}\"") } else { a }
let a = a.to_string_lossy().to_string(); }).collect();
if a.contains(' ') {
format!("\"{a}\"")
} else {
a
}
})
.collect();
format!("{prog} {}", args.join(" ")) format!("{prog} {}", args.join(" "))
} }
fn combo(ui: &mut egui::Ui, id: &str, options: &[&str], val: &mut String) { fn combo(ui: &mut egui::Ui, id: &str, options: &[&str], val: &mut String) {
egui::ComboBox::from_id_salt(id) egui::ComboBox::from_id_salt(id).selected_text(val.clone()).show_ui(ui, |ui| {
.selected_text(val.clone()) for o in options {
.show_ui(ui, |ui| { ui.selectable_value(val, (*o).to_string(), *o);
for o in options { }
ui.selectable_value(val, (*o).to_string(), *o); });
}
});
} }
impl eframe::App for App { impl eframe::App for App {
@@ -490,11 +382,7 @@ impl eframe::App for App {
h.poll(&mut self.log); h.poll(&mut self.log);
if let Some(code) = h.code { if let Some(code) = h.code {
if !self.status.starts_with("Cancelled") { if !self.status.starts_with("Cancelled") {
self.status = if code == 0 { self.status = if code == 0 { "Pipeline complete.".into() } else { format!("Pipeline exited with code {code}. See log.") };
"Pipeline complete.".into()
} else {
format!("Pipeline exited with code {code}. See log.")
};
} }
} }
} }
@@ -508,28 +396,18 @@ impl eframe::App for App {
let status = if self.is_running() { let status = if self.is_running() {
if let Some(h) = self.run.as_ref() { if let Some(h) = self.run.as_ref() {
format!("{} ({}s)", self.status, h.elapsed().as_secs()) format!("{} ({}s)", self.status, h.elapsed().as_secs())
} else { } else { self.status.clone() }
self.status.clone() } else { self.status.clone() };
}
} else {
self.status.clone()
};
ui.label(egui::RichText::new(status).strong()); ui.label(egui::RichText::new(status).strong());
ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| { ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
let running = self.is_running(); let running = self.is_running();
if running { if running {
ui.spinner(); ui.spinner();
} }
if ui if ui.add_enabled(running, egui::Button::new("⏹ Stop")).clicked() {
.add_enabled(running, egui::Button::new("⏹ Stop"))
.clicked()
{
self.cancel(); self.cancel();
} }
if ui if ui.add_enabled(!running, egui::Button::new("▶ Run pipeline")).clicked() {
.add_enabled(!running, egui::Button::new("▶ Run pipeline"))
.clicked()
{
self.start(); self.start();
} }
}); });
@@ -565,11 +443,7 @@ impl eframe::App for App {
egui::ScrollArea::vertical().show(ui, |ui| { egui::ScrollArea::vertical().show(ui, |ui| {
ui.heading("360° -> COLMAP cubemap rig"); ui.heading("360° -> COLMAP cubemap rig");
ui.add_space(2.0); ui.add_space(2.0);
ui.label( ui.label(egui::RichText::new(format!("CLI: {}", self.cli_binary_str)).small().color(egui::Color32::GRAY));
egui::RichText::new(format!("CLI: {}", self.cli_binary_str))
.small()
.color(egui::Color32::GRAY),
);
ui.add_space(8.0); ui.add_space(8.0);
ui.label("Project directory"); ui.label("Project directory");
@@ -612,7 +486,7 @@ impl eframe::App for App {
egui::ComboBox::from_id_salt("start_stage") egui::ComboBox::from_id_salt("start_stage")
.selected_text(self.start_stage.clone()) .selected_text(self.start_stage.clone())
.show_ui(ui, |ui| { .show_ui(ui, |ui| {
for s in ["extract", "cubemap", "masks", "rig"] { for s in ["extract", "masks", "cubemap", "rig"] {
ui.selectable_value(&mut self.start_stage, s.into(), s); ui.selectable_value(&mut self.start_stage, s.into(), s);
} }
}); });
@@ -630,241 +504,77 @@ impl eframe::App for App {
.default_open(true) .default_open(true)
.show(ui, |ui| { .show(ui, |ui| {
let e = &mut self.extract; let e = &mut self.extract;
egui::Grid::new("extract_grid") egui::Grid::new("extract_grid").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| {
.num_columns(2) ui.label("Chunk size"); ui.add(egui::DragValue::new(&mut e.chunk_size).range(1..=10000)); ui.end_row();
.spacing([8.0, 4.0]) ui.label("Scale width"); ui.add(egui::DragValue::new(&mut e.scale_width).range(0..=8000)); ui.end_row();
.show(ui, |ui| { ui.label("Output format"); combo(ui, "extract_fmt", &["png", "jpg"], &mut e.output_format); ui.end_row();
ui.label("Chunk size"); ui.label("JPEG quality"); ui.add(egui::DragValue::new(&mut e.jpeg_quality).range(1..=100)); ui.end_row();
ui.add(egui::DragValue::new(&mut e.chunk_size).range(1..=10000)); ui.label("Batch (frames/GPU)");ui.add(egui::DragValue::new(&mut e.batch).range(1..=64)); ui.end_row();
ui.end_row(); ui.label("Selection"); combo(ui, "extract_sel", &["chunk", "spread"], &mut e.selection); ui.end_row();
ui.label("Scale width"); ui.label("Target count"); ui.add(egui::DragValue::new(&mut e.target_count).range(0..=100000)); ui.end_row();
ui.add(egui::DragValue::new(&mut e.scale_width).range(0..=8000)); ui.label("Min gap (frames)"); ui.add(egui::DragValue::new(&mut e.min_gap).range(0..=10000)); ui.end_row();
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("Output format"); ui.label("Motion mode"); combo(ui, "extract_mm", &["frame_diff", "off"], &mut e.motion_mode); ui.end_row();
combo(ui, "extract_fmt", &["png", "jpg"], &mut e.output_format); });
ui.end_row(); ui.label(egui::RichText::new("Quality thresholds (0/255 bounds = no filter):").small());
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; let e = &mut self.extract;
egui::Grid::new("extract_q") egui::Grid::new("extract_q").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| {
.num_columns(2) ui.label("Min sharpness"); ui.add(egui::DragValue::new(&mut e.min_sharpness).speed(0.05).range(0.0..=1.0e9)); ui.end_row();
.spacing([8.0, 4.0]) ui.label("Min exposure"); ui.add(egui::DragValue::new(&mut e.min_exposure).speed(1.0).range(0.0..=255.0)); ui.end_row();
.show(ui, |ui| { 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 sharpness"); 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.add( ui.label("Min motion"); ui.add(egui::DragValue::new(&mut e.min_motion).speed(1.0).range(0.0..=255.0)); ui.end_row();
egui::DragValue::new(&mut e.min_sharpness) ui.label("Max motion"); ui.add(egui::DragValue::new(&mut e.max_motion).speed(1.0).range(0.0..=1.0e9)); ui.end_row();
.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( egui::CollapsingHeader::new("Stage 2 — Masks (YOLO segmentation; needs --features masks)")
"Stage 3 — Masks (YOLO segmentation; runs on cube faces)", .show(ui, |ui| {
) let m = &mut self.masks;
.show(ui, |ui| { egui::Grid::new("masks_grid").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| {
let m = &mut self.masks; ui.label("Classes"); ui.text_edit_singleline(&mut m.classes); ui.end_row();
egui::Grid::new("masks_grid") ui.label("Model (ONNX)"); ui.horizontal(|ui| {
.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 mut s = m.model.display().to_string();
let r = ui.text_edit_singleline(&mut s); let r = ui.text_edit_singleline(&mut s);
if r.changed() { if r.changed() { m.model = PathBuf::from(s); }
m.model = PathBuf::from(s);
}
if ui.button("…").clicked() { if ui.button("…").clicked() {
if let Some(p) = rfd::FileDialog::new() if let Some(p) = rfd::FileDialog::new().add_filter("ONNX", &["onnx"]).pick_file() { m.model = p; }
.add_filter("ONNX", &["onnx"])
.pick_file()
{
m.model = p;
}
} }
}); });
ui.end_row(); ui.end_row();
ui.label("Device"); ui.label("Device"); combo(ui, "mask_dev", &["cpu", "cuda", "rocm"], &mut m.device); ui.end_row();
combo(ui, "mask_dev", &["cpu", "rocm", "migraphx"], &mut m.device); 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.end_row(); ui.label("Imgsz"); ui.add(egui::DragValue::new(&mut m.imgsz).range(64..=4096)); ui.end_row();
ui.label("Conf"); ui.label("Dilate (px)"); ui.add(egui::DragValue::new(&mut m.dilate).range(0..=100)); ui.end_row();
ui.add( ui.label("Projection"); combo(ui, "mask_proj", &["cubemap", "equirect"], &mut m.projection); ui.end_row();
egui::DragValue::new(&mut m.conf) ui.label("Face size"); ui.add(egui::DragValue::new(&mut m.face_size).range(64..=4096)); ui.end_row();
.speed(0.01) ui.label("Skip faces"); ui.text_edit_singleline(&mut m.skip_faces); ui.end_row();
.range(0.0..=1.0) ui.label("Tile batch"); ui.add(egui::DragValue::new(&mut m.tile_batch).range(1..=64)); ui.end_row();
.fixed_decimals(2), 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.end_row();
ui.label("Imgsz");
ui.add(egui::DragValue::new(&mut m.imgsz).range(64..=4096));
ui.end_row();
ui.label("Dilate (px)");
ui.add(egui::DragValue::new(&mut m.dilate).range(0..=100));
ui.end_row();
ui.label("Skip faces");
ui.text_edit_singleline(&mut m.skip_faces);
ui.end_row();
ui.label("Batch (faces/GPU)");
ui.add(egui::DragValue::new(&mut m.batch).range(1..=64));
ui.end_row();
ui.label("NMS IoU");
ui.add(
egui::DragValue::new(&mut m.iou)
.speed(0.01)
.range(0.0..=1.0)
.fixed_decimals(2),
);
ui.end_row();
}); });
// Optional override for the MIGraphX compiled-model cache dir. ui.checkbox(&mut m.empty_if_none, "Write empty (all-keep) mask when no objects detected");
// Leave empty to use the default (~/.cache/pipeline_360/migraphx).
let mut cache_s = m
.ep_cache
.as_ref()
.map(|p| p.display().to_string())
.unwrap_or_default();
ui.horizontal(|ui| {
ui.label("GPU model cache dir:");
let r = ui.text_edit_singleline(&mut cache_s);
if r.changed() {
m.ep_cache = if cache_s.trim().is_empty() {
None
} else {
Some(PathBuf::from(cache_s.trim()))
};
}
}); });
ui.checkbox(
&mut m.empty_if_none,
"Write empty (all-keep) mask when no objects detected",
);
});
egui::CollapsingHeader::new("Stage 2 — Cubemap (equirect -> 6 faces)").show( egui::CollapsingHeader::new("Stage 3 — Cubemap (equirect -> 6 faces)")
ui, .show(ui, |ui| {
|ui| {
let c = &mut self.cubemap; let c = &mut self.cubemap;
egui::Grid::new("cubemap_grid") egui::Grid::new("cubemap_grid").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| {
.num_columns(2) ui.label("Face size"); ui.add(egui::DragValue::new(&mut c.face_size).range(64..=4096)); ui.end_row();
.spacing([8.0, 4.0]) ui.label("Skip faces"); ui.text_edit_singleline(&mut c.skip_faces); ui.end_row();
.show(ui, |ui| { ui.label("JPEG quality"); ui.add(egui::DragValue::new(&mut c.jpeg_quality).range(1..=100)); ui.end_row();
ui.label("Face size"); ui.label("Batch (frames/GPU)"); ui.add(egui::DragValue::new(&mut c.batch).range(1..=32)); ui.end_row();
ui.add(egui::DragValue::new(&mut c.face_size).range(64..=4096));
ui.end_row();
ui.label("Skip faces");
ui.text_edit_singleline(&mut c.skip_faces);
ui.end_row();
ui.label("JPEG quality");
ui.add(egui::DragValue::new(&mut c.jpeg_quality).range(1..=100));
ui.end_row();
ui.label("Batch (frames/GPU)");
ui.add(egui::DragValue::new(&mut c.batch).range(1..=32));
ui.end_row();
});
},
);
egui::CollapsingHeader::new("Stage 4 — Rig config").show(ui, |ui| {
let r = &mut self.rig;
egui::Grid::new("rig_grid")
.num_columns(2)
.spacing([8.0, 4.0])
.show(ui, |ui| {
ui.label("Reference face");
combo(
ui,
"rig_ref",
&["F", "R", "B", "L", "U", "D"],
&mut r.ref_face,
);
ui.end_row();
ui.label("Intrinsics");
combo(
ui,
"rig_int",
&["linspace", "half"],
&mut r.intrinsics_convention,
);
ui.end_row();
ui.label("Link mode");
combo(
ui,
"rig_link",
&["symlink", "hardlink", "copy"],
&mut r.link_mode,
);
ui.end_row();
}); });
}); });
egui::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();
});
});
}); });
}); });
+114 -421
View File
@@ -40,9 +40,7 @@ impl Selection {
match s { match s {
"chunk" => Ok(Selection::Chunk), "chunk" => Ok(Selection::Chunk),
"spread" => Ok(Selection::Spread), "spread" => Ok(Selection::Spread),
other => Err(anyhow!( other => Err(anyhow!("unknown --extract-selection `{other}` (chunk|spread)")),
"unknown --extract-selection `{other}` (chunk|spread)"
)),
} }
} }
} }
@@ -64,9 +62,7 @@ impl BadFrameAction {
"skip" => Ok(BadFrameAction::Skip), "skip" => Ok(BadFrameAction::Skip),
"keep" => Ok(BadFrameAction::Keep), "keep" => Ok(BadFrameAction::Keep),
"warn" => Ok(BadFrameAction::Warn), "warn" => Ok(BadFrameAction::Warn),
other => Err(anyhow!( other => Err(anyhow!("unknown --extract-bad-frame-action `{other}` (skip|keep|warn)")),
"unknown --extract-bad-frame-action `{other}` (skip|keep|warn)"
)),
} }
} }
} }
@@ -84,9 +80,7 @@ impl MotionMode {
match s { match s {
"off" => Ok(MotionMode::Off), "off" => Ok(MotionMode::Off),
"frame_diff" => Ok(MotionMode::FrameDiff), "frame_diff" => Ok(MotionMode::FrameDiff),
other => Err(anyhow!( other => Err(anyhow!("unknown --extract-motion-mode `{other}` (off|frame_diff)")),
"unknown --extract-motion-mode `{other}` (off|frame_diff)"
)),
} }
} }
} }
@@ -168,16 +162,6 @@ struct AnalyzeParams {
pad1: u32, pad1: u32,
} }
#[repr(C)]
#[derive(Clone, Copy)]
struct ThumbParams {
w: u32,
h: u32,
tw: u32,
th: u32,
pad0: u32,
}
/// Owns the wgpu pipeline + a pool of reusable resources for *batched* /// Owns the wgpu pipeline + a pool of reusable resources for *batched*
/// per-frame analysis. `batch` frames are uploaded, dispatched in one encoder /// per-frame analysis. `batch` frames are uploaded, dispatched in one encoder
/// and read back with a single `poll(Wait)`, amortising readback latency and /// and read back with a single `poll(Wait)`, amortising readback latency and
@@ -197,14 +181,6 @@ struct Analyzer<'g> {
width: u32, width: u32,
height: u32, height: u32,
batch: usize, batch: usize,
/// GPU motion-thumbnail pass (luma `THUMB_W x THUMB_H` per frame).
thumb_pipeline: wgpu::ComputePipeline,
#[allow(dead_code)]
thumb_bgl: wgpu::BindGroupLayout,
#[allow(dead_code)]
thumb_uniform: wgpu::Buffer,
thumb_bufs: Vec<wgpu::Buffer>,
thumb_bind_groups: Vec<wgpu::BindGroup>,
} }
fn pod_bytes<T: Copy>(x: &T) -> &[u8] { fn pod_bytes<T: Copy>(x: &T) -> &[u8] {
@@ -240,259 +216,98 @@ fn reduce_partial(raw: &[u8]) -> (f64, f64, f64) {
impl<'g> Analyzer<'g> { impl<'g> Analyzer<'g> {
fn new(gpu: &'g GpuContext, width: u32, height: u32, batch: usize) -> Result<Self> { fn new(gpu: &'g GpuContext, width: u32, height: u32, batch: usize) -> Result<Self> {
let batch = batch.max(1); let batch = batch.max(1);
let bgl = gpu let bgl = gpu.device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
.device label: Some("analyze-bgl"),
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { entries: &[
label: Some("analyze-bgl"), wgpu::BindGroupLayoutEntry {
entries: &[ binding: 0, visibility: wgpu::ShaderStages::COMPUTE,
wgpu::BindGroupLayoutEntry { ty: wgpu::BindingType::Buffer {
binding: 0, ty: wgpu::BufferBindingType::Uniform, has_dynamic_offset: false, min_binding_size: None,
visibility: wgpu::ShaderStages::COMPUTE, }, count: None,
ty: wgpu::BindingType::Buffer { },
ty: wgpu::BufferBindingType::Uniform, wgpu::BindGroupLayoutEntry {
has_dynamic_offset: false, binding: 1, visibility: wgpu::ShaderStages::COMPUTE,
min_binding_size: None, ty: wgpu::BindingType::Texture {
}, sample_type: wgpu::TextureSampleType::Float { filterable: false },
count: None, view_dimension: wgpu::TextureViewDimension::D2, multisampled: false,
}, }, count: None,
wgpu::BindGroupLayoutEntry { },
binding: 1, wgpu::BindGroupLayoutEntry {
visibility: wgpu::ShaderStages::COMPUTE, binding: 2, visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Texture { ty: wgpu::BindingType::Buffer {
sample_type: wgpu::TextureSampleType::Float { filterable: false }, ty: wgpu::BufferBindingType::Storage { read_only: false },
view_dimension: wgpu::TextureViewDimension::D2, has_dynamic_offset: false, min_binding_size: None,
multisampled: false, }, count: None,
}, },
count: None, ],
}, });
wgpu::BindGroupLayoutEntry { let layout = gpu.device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
binding: 2, label: Some("analyze-layout"),
visibility: wgpu::ShaderStages::COMPUTE, bind_group_layouts: &[&bgl],
ty: wgpu::BindingType::Buffer { push_constant_ranges: &[],
ty: wgpu::BufferBindingType::Storage { read_only: false }, });
has_dynamic_offset: false, let pipeline = gpu.compute_pipeline("analyze-pipeline", crate::gpu::shaders::SHARPNESS, &layout);
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_x = width.div_ceil(16);
let n_wg_y = height.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 n_wg = (n_wg_x as u64) * (n_wg_y as u64);
let partial_size = n_wg * 5 * 4; let partial_size = n_wg * 5 * 4;
let uniform = gpu.create_uniform_buffer( let uniform = gpu.create_uniform_buffer(
"analyze-params", "analyze-params",
pod_bytes(&AnalyzeParams { pod_bytes(&AnalyzeParams { width, height, pad0: 0, pad1: 0 }),
width,
height,
pad0: 0,
pad1: 0,
}),
); );
let mut src_texs = Vec::with_capacity(batch); let mut src_texs = Vec::with_capacity(batch);
let mut partials = Vec::with_capacity(batch); let mut partials = Vec::with_capacity(batch);
let mut bind_groups = Vec::with_capacity(batch); let mut bind_groups = Vec::with_capacity(batch);
for i in 0..batch { for i in 0..batch {
let tex = gpu.create_texture_2d(&format!("analyze-src-{i}"), width, height); let tex = gpu.create_texture_2d(&format!("analyze-src-{i}"), width, height);
let partial = let partial = gpu.create_storage_buffer(&format!("analyze-partial-{i}"), partial_size, None);
gpu.create_storage_buffer(&format!("analyze-partial-{i}"), partial_size, None);
let view = tex.create_view(&wgpu::TextureViewDescriptor::default()); let view = tex.create_view(&wgpu::TextureViewDescriptor::default());
let bg = gpu.device.create_bind_group(&wgpu::BindGroupDescriptor { let bg = gpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some(&format!("analyze-bg-{i}")), label: Some(&format!("analyze-bg-{i}")),
layout: &bgl, layout: &bgl,
entries: &[ entries: &[
wgpu::BindGroupEntry { wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() },
binding: 0, wgpu::BindGroupEntry { binding: 1, resource: wgpu::BindingResource::TextureView(&view) },
resource: uniform.as_entire_binding(), wgpu::BindGroupEntry { binding: 2, resource: partial.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); src_texs.push(tex);
partials.push(partial); partials.push(partial);
bind_groups.push(bg); bind_groups.push(bg);
} }
// Motion-thumbnail pass: one thread per thumbnail pixel writes its luma Ok(Analyzer { gpu, pipeline, bgl, uniform, src_texs, partials, bind_groups, n_wg_x, n_wg_y, width, height, batch })
// into a per-frame storage buffer (motion/parallax proxy on the GPU).
let thumb_bgl = gpu
.device
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("thumb-bgl"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 1,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: false },
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 2,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: false },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
],
});
let thumb_layout = gpu
.device
.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("thumb-layout"),
bind_group_layouts: &[&thumb_bgl],
push_constant_ranges: &[],
});
let thumb_pipeline = gpu.compute_pipeline(
"thumb-pipeline",
crate::gpu::shaders::THUMBNAIL,
&thumb_layout,
);
let thumb_uniform = gpu.create_uniform_buffer(
"thumb-params",
pod_bytes(&ThumbParams {
w: width,
h: height,
tw: THUMB_W,
th: THUMB_H,
pad0: 0,
}),
);
let thumb_size = (THUMB_W * THUMB_H * 4) as u64;
let mut thumb_bufs = Vec::with_capacity(batch);
let mut thumb_bind_groups = Vec::with_capacity(batch);
for i in 0..batch {
let tb = gpu.create_storage_buffer(&format!("thumb-{i}"), thumb_size, None);
let view = src_texs[i].create_view(&wgpu::TextureViewDescriptor::default());
let bg = gpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some(&format!("thumb-bg-{i}")),
layout: &thumb_bgl,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: thumb_uniform.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: wgpu::BindingResource::TextureView(&view),
},
wgpu::BindGroupEntry {
binding: 2,
resource: tb.as_entire_binding(),
},
],
});
thumb_bufs.push(tb);
thumb_bind_groups.push(bg);
}
Ok(Analyzer {
gpu,
pipeline,
bgl,
uniform,
src_texs,
partials,
bind_groups,
n_wg_x,
n_wg_y,
width,
height,
batch,
thumb_pipeline,
thumb_bgl,
thumb_uniform,
thumb_bufs,
thumb_bind_groups,
})
} }
/// Upload + dispatch a batch of frames (len <= `batch`) in one encoder. /// Upload + dispatch a batch of frames (len <= `batch`) in one encoder.
/// Runs both the sharpness pass and the motion-thumbnail pass per frame.
fn submit_batch(&self, frames: &[&[u8]]) -> Result<()> { fn submit_batch(&self, frames: &[&[u8]]) -> Result<()> {
let n = frames.len(); let n = frames.len();
debug_assert!(n <= self.batch); debug_assert!(n <= self.batch);
// Upload all frames first (separate queue writes), then encode dispatches. // Upload all frames first (separate queue writes), then encode dispatches.
for (i, rgba) in frames.iter().enumerate() { for (i, rgba) in frames.iter().enumerate() {
self.gpu self.gpu.upload_rgba(&self.src_texs[i], self.width, self.height, rgba);
.upload_rgba(&self.src_texs[i], self.width, self.height, rgba);
} }
let mut encoder = self let mut encoder = self.gpu.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
.gpu
.device
.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
{ {
let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
label: Some("analyze-pass"), label: Some("analyze-pass"), timestamp_writes: None,
timestamp_writes: None,
}); });
pass.set_pipeline(&self.pipeline); pass.set_pipeline(&self.pipeline);
for i in 0..n { for i in 0..n {
pass.set_bind_group(0, &self.bind_groups[i], &[]); pass.set_bind_group(0, &self.bind_groups[i], &[]);
pass.dispatch_workgroups(self.n_wg_x, self.n_wg_y, 1); pass.dispatch_workgroups(self.n_wg_x, self.n_wg_y, 1);
} }
// Motion-thumbnail pass (read-only src, write luma thumbnail).
pass.set_pipeline(&self.thumb_pipeline);
let thumb_wgs = (THUMB_W * THUMB_H).div_ceil(256);
for i in 0..n {
pass.set_bind_group(0, &self.thumb_bind_groups[i], &[]);
pass.dispatch_workgroups(thumb_wgs, 1, 1);
}
} }
self.gpu.queue.submit(Some(encoder.finish())); self.gpu.queue.submit(Some(encoder.finish()));
Ok(()) Ok(())
} }
/// Read back a batch of `n` partial buffers + thumbnails with one poll. /// Read back a batch of `n` partial buffers with one poll and reduce.
/// Returns `(sharpness/exposure/contrast, thumbnails)` per frame. fn collect_batch(&self, n: usize) -> Result<Vec<(f64, f64, f64)>> {
fn collect_batch(&self, n: usize) -> Result<(Vec<(f64, f64, f64)>, Vec<Vec<u8>>)> { let size = (self.n_wg_x as u64) * (self.n_wg_y as u64) * 5 * 4;
let psize = (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 mut srcs: Vec<&wgpu::Buffer> = self.partials[..n].iter().collect(); let sizes = vec![size; n];
let mut sizes: Vec<u64> = vec![psize; n];
let tsize = (THUMB_W * THUMB_H * 4) as u64;
srcs.extend(self.thumb_bufs[..n].iter());
sizes.extend(vec![tsize; n]);
let raws = self.gpu.readback_many(&srcs, &sizes)?; let raws = self.gpu.readback_many(&srcs, &sizes)?;
let mut metrics = Vec::with_capacity(n); Ok(raws.iter().map(|r| reduce_partial(r)).collect())
let mut thumbs = Vec::with_capacity(n);
for r in &raws[..n] {
metrics.push(reduce_partial(r));
}
for r in &raws[n..] {
thumbs.push(r.clone());
}
Ok((metrics, thumbs))
} }
} }
@@ -503,6 +318,24 @@ impl<'g> Analyzer<'g> {
const THUMB_W: u32 = 64; const THUMB_W: u32 = 64;
const THUMB_H: u32 = 32; 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. /// Mean absolute luminance difference between two thumbnails, in 0..255.
fn frame_motion(a: &[u8], b: &[u8]) -> f64 { fn frame_motion(a: &[u8], b: &[u8]) -> f64 {
if a.len() != b.len() { if a.len() != b.len() {
@@ -544,30 +377,12 @@ fn read_metadata(path: &Path) -> Result<Vec<FrameMetrics>> {
} }
let mut it = line.split(','); let mut it = line.split(',');
let frame = it.next().and_then(|s| s.trim().parse::<u64>().ok()); let frame = it.next().and_then(|s| s.trim().parse::<u64>().ok());
let sharp = it let sharp = it.next().and_then(|s| s.trim().parse::<f64>().ok()).unwrap_or(0.0);
.next() let exposure = it.next().and_then(|s| s.trim().parse::<f64>().ok()).unwrap_or(0.0);
.and_then(|s| s.trim().parse::<f64>().ok()) let contrast = it.next().and_then(|s| s.trim().parse::<f64>().ok()).unwrap_or(0.0);
.unwrap_or(0.0); let motion = 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 { if let Some(frame) = frame {
out.push(FrameMetrics { out.push(FrameMetrics { frame, sharpness: sharp, exposure, contrast, motion });
frame,
sharpness: sharp,
exposure,
contrast,
motion,
});
} }
} }
Ok(out) Ok(out)
@@ -616,8 +431,7 @@ fn select_frames(cfg: &ExtractConfig, metrics: &[FrameMetrics]) -> Vec<u64> {
let end = (i + chunk).min(metrics.len()); let end = (i + chunk).min(metrics.len());
let window = &metrics[i..end]; let window = &metrics[i..end];
// Prefer the sharpest surviving frame; fall back per bad-frame-action. // Prefer the sharpest surviving frame; fall back per bad-frame-action.
let survivors: Vec<&FrameMetrics> = let survivors: Vec<&FrameMetrics> = window.iter().filter(|m| passes(cfg, m)).collect();
window.iter().filter(|m| passes(cfg, m)).collect();
let pick = survivors let pick = survivors
.iter() .iter()
.max_by(|a, b| a.sharpness.partial_cmp(&b.sharpness).unwrap()) .max_by(|a, b| a.sharpness.partial_cmp(&b.sharpness).unwrap())
@@ -626,10 +440,12 @@ fn select_frames(cfg: &ExtractConfig, metrics: &[FrameMetrics]) -> Vec<u64> {
Some(m) => Some(m.frame), Some(m) => Some(m.frame),
None => match cfg.bad_frame_action { None => match cfg.bad_frame_action {
BadFrameAction::Skip => None, BadFrameAction::Skip => None,
BadFrameAction::Warn | BadFrameAction::Keep => window BadFrameAction::Warn | BadFrameAction::Keep => {
.iter() window
.max_by(|a, b| a.sharpness.partial_cmp(&b.sharpness).unwrap()) .iter()
.map(|m| m.frame), .max_by(|a, b| a.sharpness.partial_cmp(&b.sharpness).unwrap())
.map(|m| m.frame)
}
}, },
}; };
if let Some(f) = chosen { if let Some(f) = chosen {
@@ -649,11 +465,7 @@ fn select_frames(cfg: &ExtractConfig, metrics: &[FrameMetrics]) -> Vec<u64> {
let mut survivors: Vec<&FrameMetrics> = let mut survivors: Vec<&FrameMetrics> =
metrics.iter().filter(|m| passes(cfg, m)).collect(); metrics.iter().filter(|m| passes(cfg, m)).collect();
survivors.sort_by(|a, b| b.sharpness.partial_cmp(&a.sharpness).unwrap()); survivors.sort_by(|a, b| b.sharpness.partial_cmp(&a.sharpness).unwrap());
let target = if cfg.target_count == 0 { let target = if cfg.target_count == 0 { survivors.len() as u32 } else { cfg.target_count };
survivors.len() as u32
} else {
cfg.target_count
};
let mut accepted: Vec<u64> = Vec::new(); let mut accepted: Vec<u64> = Vec::new();
for m in survivors { for m in survivors {
if accepted.len() as u32 >= target { if accepted.len() as u32 >= target {
@@ -678,12 +490,7 @@ fn select_frames(cfg: &ExtractConfig, metrics: &[FrameMetrics]) -> Vec<u64> {
// =========================================================================== // ===========================================================================
/// Format a printf-style pattern with one integer (supports `%d`, `%05d`, …). /// Format a printf-style pattern with one integer (supports `%d`, `%05d`, …).
pub fn format_output_filename( pub fn format_output_filename(pattern: &str, index: u32, output_format: &str, prefix: &str) -> Result<String> {
pattern: &str,
index: u32,
output_format: &str,
prefix: &str,
) -> Result<String> {
// Find the single `%[flags][width]d` token and substitute. // Find the single `%[flags][width]d` token and substitute.
let bytes = pattern.as_bytes(); let bytes = pattern.as_bytes();
let mut out = String::new(); let mut out = String::new();
@@ -714,21 +521,9 @@ pub fn format_output_filename(
let composed = format!("{prefix}{out}"); let composed = format!("{prefix}{out}");
let path = Path::new(&composed); let path = Path::new(&composed);
let target_suffix = format!(".{output_format}"); let target_suffix = format!(".{output_format}");
let final_name = if path let final_name = if path.extension().map(|e| e.to_string_lossy().to_lowercase()).as_deref() == Some("jpg")
.extension() || path.extension().map(|e| e.to_string_lossy().to_lowercase()).as_deref() == Some("jpeg")
.map(|e| e.to_string_lossy().to_lowercase()) || path.extension().map(|e| e.to_string_lossy().to_lowercase()).as_deref() == Some("png")
.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('.')) path.with_extension(target_suffix.trim_start_matches('.'))
.to_string_lossy() .to_string_lossy()
@@ -745,14 +540,7 @@ pub fn format_output_filename(
// Frame writing // Frame writing
// =========================================================================== // ===========================================================================
fn save_frame( fn save_frame(rgba: &[u8], w: u32, h: u32, path: &Path, jpeg_quality: u32, output_format: &str) -> Result<()> {
rgba: &[u8],
w: u32,
h: u32,
path: &Path,
jpeg_quality: u32,
output_format: &str,
) -> Result<()> {
if let Some(parent) = path.parent() { if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?; std::fs::create_dir_all(parent)?;
} }
@@ -764,10 +552,8 @@ fn save_frame(
}; };
if fmt == image::ImageFormat::Jpeg { if fmt == image::ImageFormat::Jpeg {
let rgb = image::DynamicImage::ImageRgba8(img).to_rgb8(); let rgb = image::DynamicImage::ImageRgba8(img).to_rgb8();
let mut out = std::fs::File::create(path) let mut out = std::fs::File::create(path).with_context(|| format!("cannot create {}", path.display()))?;
.with_context(|| format!("cannot create {}", path.display()))?; let mut enc = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut out, jpeg_quality as u8);
let mut enc =
image::codecs::jpeg::JpegEncoder::new_with_quality(&mut out, jpeg_quality as u8);
enc.encode_image(&image::DynamicImage::ImageRgb8(rgb)) enc.encode_image(&image::DynamicImage::ImageRgb8(rgb))
.with_context(|| format!("cannot encode jpeg {}", path.display()))?; .with_context(|| format!("cannot encode jpeg {}", path.display()))?;
} else { } else {
@@ -801,7 +587,7 @@ fn analyze_video(
) -> Result<Vec<FrameMetrics>> { ) -> Result<Vec<FrameMetrics>> {
let (w, h) = video::scaled_dims(probe, cfg.scale_width); let (w, h) = video::scaled_dims(probe, cfg.scale_width);
logging::emit( logging::emit(
&format!("Analyzing sharpness @ {w}x{h} (GPU Laplacian + luma stats)…",), &format!("Analyzing sharpness @ {w}x{h} (GPU Laplacian + luma stats)…", ),
logger, logger,
); );
let batch = cfg.batch.max(1) as usize; let batch = cfg.batch.max(1) as usize;
@@ -828,32 +614,18 @@ fn analyze_video(
} }
let refs: Vec<&[u8]> = buf.iter().map(|v| v.as_slice()).collect(); let refs: Vec<&[u8]> = buf.iter().map(|v| v.as_slice()).collect();
analyzer.submit_batch(&refs)?; analyzer.submit_batch(&refs)?;
let (metrics_b, thumbs_b) = analyzer.collect_batch(refs.len())?; let results = analyzer.collect_batch(refs.len())?;
for (i, (sharp, exp, contrast)) in metrics_b.into_iter().enumerate() { for (i, (sharp, exp, contrast)) in results.into_iter().enumerate() {
let rgba = &buf[i];
let motion = if want_motion { let motion = if want_motion {
// Thumbnails come back from the GPU as an `array<f32>` of luma let thumb = thumbnail(rgba, w, h);
// (0..255). Decode each 4-byte element to an f32, then round to let m = prev_thumb.as_ref().map(|p| frame_motion(&thumb, p)).unwrap_or(0.0);
// u8 for the cheap CPU motion diff. prev_thumb = Some(thumb);
let thumb_u8: Vec<u8> = thumbs_b[i]
.chunks_exact(4)
.map(|px| f32::from_le_bytes([px[0], px[1], px[2], px[3]]).round() as u8)
.collect();
let m = prev_thumb
.as_ref()
.map(|p| frame_motion(&thumb_u8, p))
.unwrap_or(0.0);
prev_thumb = Some(thumb_u8);
m m
} else { } else {
0.0 0.0
}; };
metrics.push(FrameMetrics { metrics.push(FrameMetrics { frame: frame_no, sharpness: sharp, exposure: exp, contrast, motion });
frame: frame_no,
sharpness: sharp,
exposure: exp,
contrast,
motion,
});
frame_no += 1; frame_no += 1;
} }
if frame_no.is_multiple_of(50) || buf.len() < batch { if frame_no.is_multiple_of(50) || buf.len() < batch {
@@ -888,12 +660,7 @@ fn extract_frames(
std::fs::create_dir_all(output_dir)?; std::fs::create_dir_all(output_dir)?;
let targets: std::collections::BTreeSet<u64> = selected.iter().copied().collect(); let targets: std::collections::BTreeSet<u64> = selected.iter().copied().collect();
logging::emit( logging::emit(
&format!( &format!("Extracting {} frames at {}x{}…", targets.len(), probe.width, probe.height),
"Extracting {} frames at {}x{}…",
targets.len(),
probe.width,
probe.height
),
logger, logger,
); );
let mut pipe = FramePipe::open(path, 0, probe)?; let mut pipe = FramePipe::open(path, 0, probe)?;
@@ -903,20 +670,8 @@ fn extract_frames(
while let Some(rgba) = pipe.next_frame()? { while let Some(rgba) = pipe.next_frame()? {
should_cancel.check().map_err(|_| anyhow!("Cancelled."))?; should_cancel.check().map_err(|_| anyhow!("Cancelled."))?;
if targets.contains(&frame_no) { if targets.contains(&frame_no) {
let name = format_output_filename( let name = format_output_filename(&cfg.output_pattern, output_index, &cfg.output_format, prefix)?;
&cfg.output_pattern, save_frame(&rgba, pipe.width, pipe.height, &output_dir.join(name), cfg.jpeg_quality, &cfg.output_format)?;
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; output_index += 1;
saved += 1; saved += 1;
logging::emit_progress(&format!("extracted {saved}/{}", targets.len()), logger); logging::emit_progress(&format!("extracted {saved}/{}", targets.len()), logger);
@@ -928,11 +683,7 @@ fn extract_frames(
} }
pipe.close(); pipe.close();
if saved != targets.len() { if saved != targets.len() {
bail!( bail!("Some selected frames could not be extracted. Expected {}, saved {}.", targets.len(), saved);
"Some selected frames could not be extracted. Expected {}, saved {}.",
targets.len(),
saved
);
} }
Ok(saved) Ok(saved)
} }
@@ -942,17 +693,10 @@ fn per_video_prefixes(videos: &[PathBuf]) -> Vec<String> {
let mut used: std::collections::HashMap<String, u32> = std::collections::HashMap::new(); let mut used: std::collections::HashMap<String, u32> = std::collections::HashMap::new();
let mut out = Vec::with_capacity(videos.len()); let mut out = Vec::with_capacity(videos.len());
for v in videos { for v in videos {
let base = v let base = v.file_stem().map(|s| s.to_string_lossy().to_string()).unwrap_or_default();
.file_stem()
.map(|s| s.to_string_lossy().to_string())
.unwrap_or_default();
let n = used.entry(base.clone()).or_insert(0); let n = used.entry(base.clone()).or_insert(0);
*n += 1; *n += 1;
out.push(if *n == 1 { out.push(if *n == 1 { base } else { format!("{base}_{}", *n) });
base
} else {
format!("{base}_{}", *n)
});
} }
out out
} }
@@ -989,15 +733,7 @@ pub fn run_extract_stage(
for (idx, (video, prefix)) in videos.iter().zip(prefixes.iter()).enumerate() { for (idx, (video, prefix)) in videos.iter().zip(prefixes.iter()).enumerate() {
should_cancel.check().map_err(|_| anyhow!("Cancelled."))?; should_cancel.check().map_err(|_| anyhow!("Cancelled."))?;
if multi { if multi {
logging::emit( logging::emit(&format!("--- Source {}/{}: {} ---", idx + 1, videos.len(), video.display()), logger);
&format!(
"--- Source {}/{}: {} ---",
idx + 1,
videos.len(),
video.display()
),
logger,
);
} }
let tag = if multi { prefix.as_str() } else { "" }; let tag = if multi { prefix.as_str() } else { "" };
let metadata_path = output_dir.join(if multi { let metadata_path = output_dir.join(if multi {
@@ -1008,36 +744,18 @@ pub fn run_extract_stage(
let probe = video::probe(video)?; let probe = video::probe(video)?;
logging::emit( logging::emit(
&format!( &format!("Video: {}x{} {} fps, {} frames declared",
"Video: {}x{} {} fps, {} frames declared", probe.width, probe.height,
probe.width, probe.fps_f64().map(|f| format!("{f:.3}")).unwrap_or_else(|| "?".into()),
probe.height, probe.nb_frames.map(|n| n.to_string()).unwrap_or_else(|| "?".into())),
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, logger,
); );
let metrics = if metadata_path.exists() && cfg.reuse_metadata { let metrics = if metadata_path.exists() && cfg.reuse_metadata {
logging::emit( logging::emit(&format!("Using existing sharpness metadata: {}", metadata_path.display()), logger);
&format!(
"Using existing sharpness metadata: {}",
metadata_path.display()
),
logger,
);
read_metadata(&metadata_path)? read_metadata(&metadata_path)?
} else { } else {
logging::emit( logging::emit(&format!("Generating sharpness metadata: {}", metadata_path.display()), logger);
&format!("Generating sharpness metadata: {}", metadata_path.display()),
logger,
);
let m = analyze_video(gpu, video, &probe, cfg, should_cancel, logger)?; let m = analyze_video(gpu, video, &probe, cfg, should_cancel, logger)?;
write_metadata(&metadata_path, &m)?; write_metadata(&metadata_path, &m)?;
m m
@@ -1048,46 +766,21 @@ pub fn run_extract_stage(
let n_rejected = metrics.len() - n_survivors; let n_rejected = metrics.len() - n_survivors;
if n_rejected > 0 { if n_rejected > 0 {
logging::emit( logging::emit(
&format!( &format!("Quality filter: {n_rejected}/{} frames rejected (thresholds).", metrics.len()),
"Quality filter: {n_rejected}/{} frames rejected (thresholds).",
metrics.len()
),
logger, logger,
); );
} }
logging::emit( logging::emit(
&format!( &format!("Selected {}/{} frames (mode={:?}, chunk={}, gap={}).",
"Selected {}/{} frames (mode={:?}, chunk={}, gap={}).", selected.len(), metrics.len(), cfg.selection, cfg.chunk_size, cfg.min_gap),
selected.len(),
metrics.len(),
cfg.selection,
cfg.chunk_size,
cfg.min_gap
),
logger, logger,
); );
let saved = extract_frames( let saved = extract_frames(video, &probe, &selected, output_dir, cfg, tag, should_cancel, logger)?;
video, outcomes.push(ExtractOutcome { video: video.to_path_buf(), prefix: prefix.clone(), frames_written: saved });
&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("Done: sharp frames extracted.", logger);
logging::emit( logging::emit(&format!("Output directory: {}", output_dir.display()), logger);
&format!("Output directory: {}", output_dir.display()),
logger,
);
Ok(outcomes) Ok(outcomes)
} }