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