This commit is contained in:
2026-08-22 23:02:46 +02:00
parent 24a546acba
commit f918231146
3 changed files with 846 additions and 247 deletions
+2
View File
@@ -463,6 +463,8 @@ 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");
} }
+423 -133
View File
@@ -68,16 +68,18 @@ 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 let mut child = cmd.stdout(Stdio::piped()).stderr(Stdio::piped()).spawn()?;
.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 { child, rx, started: Instant::now(), code: None }) Ok(RunHandle {
child,
rx,
started: Instant::now(),
code: None,
})
} }
fn poll(&mut self, log: &mut VecDeque<String>) { fn poll(&mut self, log: &mut VecDeque<String>) {
@@ -110,7 +112,11 @@ impl RunHandle {
} }
} }
fn spawn_line_reader<R: std::io::Read + Send + 'static>(mut r: R, tx: Sender<Line>, is_stderr: bool) { fn spawn_line_reader<R: std::io::Read + Send + 'static>(
mut r: R,
tx: Sender<Line>,
is_stderr: bool,
) {
std::thread::spawn(move || { std::thread::spawn(move || {
let mut buf = [0u8; 8192]; let mut buf = [0u8; 8192];
let mut leftover = Vec::new(); let mut leftover = Vec::new();
@@ -167,10 +173,22 @@ struct ExtractForm {
impl Default for ExtractForm { impl Default for ExtractForm {
fn default() -> Self { fn default() -> Self {
Self { Self {
chunk_size: 30, scale_width: 1920, output_format: "png".into(), jpeg_quality: 95, chunk_size: 30,
min_sharpness: 0.0, min_exposure: 0.0, max_exposure: 255.0, min_contrast: 0.0, scale_width: 1920,
min_motion: 0.0, max_motion: 1.0e12, selection: "chunk".into(), target_count: 0, output_format: "png".into(),
min_gap: 0, bad_frame_action: "keep".into(), motion_mode: "frame_diff".into(), batch: 8, 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, conf: f64,
imgsz: u32, imgsz: u32,
dilate: u32, dilate: u32,
projection: String,
face_size: u32,
skip_faces: String, skip_faces: String,
tile_batch: u32, 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"), device: "cpu".into(), conf: 0.35, model: PathBuf::from("yolo11x-seg.onnx"),
imgsz: 1280, dilate: 6, projection: "cubemap".into(), face_size: 1024, device: "cpu".into(),
skip_faces: String::new(), tile_batch: 8, empty_if_none: false, iou: 0.45, 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)] #[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 { 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)] #[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 { 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 { fn default() -> Self {
let cli = resolve_cli_binary(); let cli = resolve_cli_binary();
Self { Self {
project_dir: PathBuf::new(), videos: Vec::new(), start_stage: "extract".into(), project_dir: PathBuf::new(),
skip_masks: false, extract: ExtractForm::default(), masks: MasksForm::default(), videos: Vec::new(),
cubemap: CubemapForm::default(), rig: RigForm::default(), start_stage: "extract".into(),
cli_binary_str: cli.as_ref().map(|p| p.display().to_string()).unwrap_or_else(|| "pipeline_360 (NOT FOUND — build it first)".into()), skip_masks: false,
cli_binary: cli, log: VecDeque::new(), run: None, status: "Idle".into(), last_progress: String::new(), 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; let e = &self.extract;
cmd.args([ cmd.args([
"--extract-chunk-size", &e.chunk_size.to_string(), "--extract-chunk-size",
"--extract-scale-width", &e.scale_width.to_string(), &e.chunk_size.to_string(),
"--extract-output-format", &e.output_format, "--extract-scale-width",
"--extract-jpeg-quality", &e.jpeg_quality.to_string(), &e.scale_width.to_string(),
"--extract-min-sharpness", &fmt_f(e.min_sharpness), "--extract-output-format",
"--extract-min-exposure", &fmt_f(e.min_exposure), &e.output_format,
"--extract-max-exposure", &fmt_f(e.max_exposure), "--extract-jpeg-quality",
"--extract-min-contrast", &fmt_f(e.min_contrast), &e.jpeg_quality.to_string(),
"--extract-min-motion", &fmt_f(e.min_motion), "--extract-min-sharpness",
"--extract-max-motion", &fmt_f(e.max_motion), &fmt_f(e.min_sharpness),
"--extract-selection", &e.selection, "--extract-min-exposure",
"--extract-target-count", &e.target_count.to_string(), &fmt_f(e.min_exposure),
"--extract-min-gap", &e.min_gap.to_string(), "--extract-max-exposure",
"--extract-bad-frame-action", &e.bad_frame_action, &fmt_f(e.max_exposure),
"--extract-motion-mode", &e.motion_mode, "--extract-min-contrast",
"--extract-batch", &e.batch.to_string(), &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; let m = &self.masks;
cmd.args([ cmd.args([
"--mask-classes", &m.classes, "--mask-classes",
"--mask-model", &m.model.display().to_string(), &m.classes,
"--mask-device", &m.device, "--mask-model",
"--mask-conf", &fmt_f(m.conf), &m.model.display().to_string(),
"--mask-imgsz", &m.imgsz.to_string(), "--mask-device",
"--mask-dilate", &m.dilate.to_string(), &m.device,
"--mask-projection", &m.projection, "--mask-conf",
"--mask-face-size", &m.face_size.to_string(), &fmt_f(m.conf),
"--mask-skip-faces", &m.skip_faces, "--mask-imgsz",
"--mask-tile-batch", &m.tile_batch.to_string(), &m.imgsz.to_string(),
"--mask-iou", &fmt_f(m.iou), "--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 { 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", &c.face_size.to_string(), "--cubemap-face-size",
"--cubemap-skip-faces", &c.skip_faces, &c.face_size.to_string(),
"--cubemap-jpeg-quality", &c.jpeg_quality.to_string(), "--cubemap-skip-faces",
"--cubemap-batch", &c.batch.to_string(), &c.skip_faces,
"--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", &r.ref_face, "--rig-ref-face",
"--rig-intrinsics-convention", &r.intrinsics_convention, &r.ref_face,
"--rig-link-mode", &r.link_mode, "--rig-intrinsics-convention",
&r.intrinsics_convention,
"--rig-link-mode",
&r.link_mode,
]); ]);
Some(cmd) Some(cmd)
} }
@@ -334,7 +427,9 @@ 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 { return }; let Some(cmd) = self.build_command() else {
return;
};
self.log.clear(); self.log.clear();
self.last_progress.clear(); self.last_progress.clear();
self.status = "Running…".into(); self.status = "Running…".into();
@@ -355,24 +450,37 @@ 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") { "1e12".into() } else { s } if s.eq_ignore_ascii_case("inf") {
"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.get_args().map(|a| { let args: Vec<String> = cmd
let a = a.to_string_lossy().to_string(); .get_args()
if a.contains(' ') { format!("\"{a}\"") } else { a } .map(|a| {
}).collect(); let a = a.to_string_lossy().to_string();
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).selected_text(val.clone()).show_ui(ui, |ui| { egui::ComboBox::from_id_salt(id)
for o in options { .selected_text(val.clone())
ui.selectable_value(val, (*o).to_string(), *o); .show_ui(ui, |ui| {
} for o in options {
}); ui.selectable_value(val, (*o).to_string(), *o);
}
});
} }
impl eframe::App for App { impl eframe::App for App {
@@ -382,7 +490,11 @@ 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 { "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() { 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 { self.status.clone() } } else {
} else { self.status.clone() }; 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.add_enabled(running, egui::Button::new("⏹ Stop")).clicked() { if ui
.add_enabled(running, egui::Button::new("⏹ Stop"))
.clicked()
{
self.cancel(); 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(); self.start();
} }
}); });
@@ -443,7 +565,11 @@ 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(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.add_space(8.0);
ui.label("Project directory"); ui.label("Project directory");
@@ -486,7 +612,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", "masks", "cubemap", "rig"] { for s in ["extract", "cubemap", "masks", "rig"] {
ui.selectable_value(&mut self.start_stage, s.into(), s); ui.selectable_value(&mut self.start_stage, s.into(), s);
} }
}); });
@@ -504,77 +630,241 @@ 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").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| { egui::Grid::new("extract_grid")
ui.label("Chunk size"); ui.add(egui::DragValue::new(&mut e.chunk_size).range(1..=10000)); ui.end_row(); .num_columns(2)
ui.label("Scale width"); ui.add(egui::DragValue::new(&mut e.scale_width).range(0..=8000)); ui.end_row(); .spacing([8.0, 4.0])
ui.label("Output format"); combo(ui, "extract_fmt", &["png", "jpg"], &mut e.output_format); ui.end_row(); .show(ui, |ui| {
ui.label("JPEG quality"); ui.add(egui::DragValue::new(&mut e.jpeg_quality).range(1..=100)); ui.end_row(); ui.label("Chunk size");
ui.label("Batch (frames/GPU)");ui.add(egui::DragValue::new(&mut e.batch).range(1..=64)); ui.end_row(); ui.add(egui::DragValue::new(&mut e.chunk_size).range(1..=10000));
ui.label("Selection"); combo(ui, "extract_sel", &["chunk", "spread"], &mut e.selection); ui.end_row(); ui.end_row();
ui.label("Target count"); ui.add(egui::DragValue::new(&mut e.target_count).range(0..=100000)); ui.end_row(); ui.label("Scale width");
ui.label("Min gap (frames)"); ui.add(egui::DragValue::new(&mut e.min_gap).range(0..=10000)); ui.end_row(); ui.add(egui::DragValue::new(&mut e.scale_width).range(0..=8000));
ui.label("Bad-frame action"); combo(ui, "extract_bfa", &["keep", "skip", "warn"], &mut e.bad_frame_action); ui.end_row(); ui.end_row();
ui.label("Motion mode"); combo(ui, "extract_mm", &["frame_diff", "off"], &mut e.motion_mode); ui.end_row(); ui.label("Output format");
}); combo(ui, "extract_fmt", &["png", "jpg"], &mut e.output_format);
ui.label(egui::RichText::new("Quality thresholds (0/255 bounds = no filter):").small()); 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; let e = &mut self.extract;
egui::Grid::new("extract_q").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| { egui::Grid::new("extract_q")
ui.label("Min sharpness"); ui.add(egui::DragValue::new(&mut e.min_sharpness).speed(0.05).range(0.0..=1.0e9)); ui.end_row(); .num_columns(2)
ui.label("Min exposure"); ui.add(egui::DragValue::new(&mut e.min_exposure).speed(1.0).range(0.0..=255.0)); ui.end_row(); .spacing([8.0, 4.0])
ui.label("Max exposure"); ui.add(egui::DragValue::new(&mut e.max_exposure).speed(1.0).range(0.0..=255.0)); ui.end_row(); .show(ui, |ui| {
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 sharpness");
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.add(
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::DragValue::new(&mut e.min_sharpness)
}); .speed(0.05)
.range(0.0..=1.0e9),
);
ui.end_row();
ui.label("Min exposure");
ui.add(
egui::DragValue::new(&mut e.min_exposure)
.speed(1.0)
.range(0.0..=255.0),
);
ui.end_row();
ui.label("Max exposure");
ui.add(
egui::DragValue::new(&mut e.max_exposure)
.speed(1.0)
.range(0.0..=255.0),
);
ui.end_row();
ui.label("Min contrast");
ui.add(
egui::DragValue::new(&mut e.min_contrast)
.speed(1.0)
.range(0.0..=255.0),
);
ui.end_row();
ui.label("Min motion");
ui.add(
egui::DragValue::new(&mut e.min_motion)
.speed(1.0)
.range(0.0..=255.0),
);
ui.end_row();
ui.label("Max motion");
ui.add(
egui::DragValue::new(&mut e.max_motion)
.speed(1.0)
.range(0.0..=1.0e9),
);
ui.end_row();
});
}); });
egui::CollapsingHeader::new("Stage 2 — Masks (YOLO segmentation; needs --features masks)") egui::CollapsingHeader::new(
.show(ui, |ui| { "Stage 3 — Masks (YOLO segmentation; runs on cube faces)",
let m = &mut self.masks; )
egui::Grid::new("masks_grid").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| { .show(ui, |ui| {
ui.label("Classes"); ui.text_edit_singleline(&mut m.classes); ui.end_row(); let m = &mut self.masks;
ui.label("Model (ONNX)"); ui.horizontal(|ui| { egui::Grid::new("masks_grid")
.num_columns(2)
.spacing([8.0, 4.0])
.show(ui, |ui| {
ui.label("Classes");
ui.text_edit_singleline(&mut m.classes);
ui.end_row();
ui.label("Model (ONNX)");
ui.horizontal(|ui| {
let mut s = m.model.display().to_string(); let 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() { m.model = PathBuf::from(s); } if r.changed() {
m.model = PathBuf::from(s);
}
if ui.button("…").clicked() { 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.end_row();
ui.label("Device"); combo(ui, "mask_dev", &["cpu", "cuda", "rocm"], &mut m.device); ui.end_row(); ui.label("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(); combo(ui, "mask_dev", &["cpu", "rocm", "migraphx"], &mut m.device);
ui.label("Imgsz"); ui.add(egui::DragValue::new(&mut m.imgsz).range(64..=4096)); ui.end_row(); ui.end_row();
ui.label("Dilate (px)"); ui.add(egui::DragValue::new(&mut m.dilate).range(0..=100)); ui.end_row(); ui.label("Conf");
ui.label("Projection"); combo(ui, "mask_proj", &["cubemap", "equirect"], &mut m.projection); ui.end_row(); ui.add(
ui.label("Face size"); ui.add(egui::DragValue::new(&mut m.face_size).range(64..=4096)); ui.end_row(); egui::DragValue::new(&mut m.conf)
ui.label("Skip faces"); ui.text_edit_singleline(&mut m.skip_faces); ui.end_row(); .speed(0.01)
ui.label("Tile batch"); ui.add(egui::DragValue::new(&mut m.tile_batch).range(1..=64)); ui.end_row(); .range(0.0..=1.0)
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(); .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)") egui::CollapsingHeader::new("Stage 2 — Cubemap (equirect -> 6 faces)").show(
.show(ui, |ui| { ui,
|ui| {
let c = &mut self.cubemap; let c = &mut self.cubemap;
egui::Grid::new("cubemap_grid").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| { egui::Grid::new("cubemap_grid")
ui.label("Face size"); ui.add(egui::DragValue::new(&mut c.face_size).range(64..=4096)); ui.end_row(); .num_columns(2)
ui.label("Skip faces"); ui.text_edit_singleline(&mut c.skip_faces); ui.end_row(); .spacing([8.0, 4.0])
ui.label("JPEG quality"); ui.add(egui::DragValue::new(&mut c.jpeg_quality).range(1..=100)); ui.end_row(); .show(ui, |ui| {
ui.label("Batch (frames/GPU)"); ui.add(egui::DragValue::new(&mut c.batch).range(1..=32)); ui.end_row(); ui.label("Face size");
}); ui.add(egui::DragValue::new(&mut c.face_size).range(64..=4096));
}); ui.end_row();
ui.label("Skip faces");
ui.text_edit_singleline(&mut c.skip_faces);
ui.end_row();
ui.label("JPEG quality");
ui.add(egui::DragValue::new(&mut c.jpeg_quality).range(1..=100));
ui.end_row();
ui.label("Batch (frames/GPU)");
ui.add(egui::DragValue::new(&mut c.batch).range(1..=32));
ui.end_row();
});
},
);
egui::CollapsingHeader::new("Stage 4 — Rig config") egui::CollapsingHeader::new("Stage 4 — Rig config").show(ui, |ui| {
.show(ui, |ui| { let r = &mut self.rig;
let r = &mut self.rig; egui::Grid::new("rig_grid")
egui::Grid::new("rig_grid").num_columns(2).spacing([8.0, 4.0]).show(ui, |ui| { .num_columns(2)
ui.label("Reference face"); combo(ui, "rig_ref", &["F", "R", "B", "L", "U", "D"], &mut r.ref_face); ui.end_row(); .spacing([8.0, 4.0])
ui.label("Intrinsics"); combo(ui, "rig_int", &["linspace", "half"], &mut r.intrinsics_convention); ui.end_row(); .show(ui, |ui| {
ui.label("Link mode"); combo(ui, "rig_link", &["symlink", "hardlink", "copy"], &mut r.link_mode); ui.end_row(); 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();
}); });
}); });
}); });
}); });
+421 -114
View File
@@ -40,7 +40,9 @@ 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!("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), "skip" => Ok(BadFrameAction::Skip),
"keep" => Ok(BadFrameAction::Keep), "keep" => Ok(BadFrameAction::Keep),
"warn" => Ok(BadFrameAction::Warn), "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 { match s {
"off" => Ok(MotionMode::Off), "off" => Ok(MotionMode::Off),
"frame_diff" => Ok(MotionMode::FrameDiff), "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, 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
@@ -181,6 +197,14 @@ 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] {
@@ -216,98 +240,259 @@ 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.device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { let bgl = gpu
label: Some("analyze-bgl"), .device
entries: &[ .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
wgpu::BindGroupLayoutEntry { label: Some("analyze-bgl"),
binding: 0, visibility: wgpu::ShaderStages::COMPUTE, entries: &[
ty: wgpu::BindingType::Buffer { wgpu::BindGroupLayoutEntry {
ty: wgpu::BufferBindingType::Uniform, has_dynamic_offset: false, min_binding_size: None, binding: 0,
}, count: None, visibility: wgpu::ShaderStages::COMPUTE,
}, ty: wgpu::BindingType::Buffer {
wgpu::BindGroupLayoutEntry { ty: wgpu::BufferBindingType::Uniform,
binding: 1, visibility: wgpu::ShaderStages::COMPUTE, has_dynamic_offset: false,
ty: wgpu::BindingType::Texture { min_binding_size: None,
sample_type: wgpu::TextureSampleType::Float { filterable: false }, },
view_dimension: wgpu::TextureViewDimension::D2, multisampled: false, count: None,
}, count: None, },
}, wgpu::BindGroupLayoutEntry {
wgpu::BindGroupLayoutEntry { binding: 1,
binding: 2, visibility: wgpu::ShaderStages::COMPUTE, visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer { ty: wgpu::BindingType::Texture {
ty: wgpu::BufferBindingType::Storage { read_only: false }, sample_type: wgpu::TextureSampleType::Float { filterable: false },
has_dynamic_offset: false, min_binding_size: None, view_dimension: wgpu::TextureViewDimension::D2,
}, count: None, multisampled: false,
}, },
], count: None,
}); },
let layout = gpu.device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { wgpu::BindGroupLayoutEntry {
label: Some("analyze-layout"), binding: 2,
bind_group_layouts: &[&bgl], visibility: wgpu::ShaderStages::COMPUTE,
push_constant_ranges: &[], ty: wgpu::BindingType::Buffer {
}); ty: wgpu::BufferBindingType::Storage { read_only: false },
let pipeline = gpu.compute_pipeline("analyze-pipeline", crate::gpu::shaders::SHARPNESS, &layout); 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_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 { 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 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 = 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 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 { binding: 0, resource: uniform.as_entire_binding() }, wgpu::BindGroupEntry {
wgpu::BindGroupEntry { binding: 1, resource: wgpu::BindingResource::TextureView(&view) }, binding: 0,
wgpu::BindGroupEntry { binding: 2, resource: partial.as_entire_binding() }, 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); src_texs.push(tex);
partials.push(partial); partials.push(partial);
bind_groups.push(bg); 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. /// 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.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 { 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); 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 with one poll and reduce. /// Read back a batch of `n` partial buffers + thumbnails with one poll.
fn collect_batch(&self, n: usize) -> Result<Vec<(f64, f64, f64)>> { /// Returns `(sharpness/exposure/contrast, thumbnails)` per frame.
let size = (self.n_wg_x as u64) * (self.n_wg_y as u64) * 5 * 4; fn collect_batch(&self, n: usize) -> Result<(Vec<(f64, f64, f64)>, Vec<Vec<u8>>)> {
let srcs: Vec<&wgpu::Buffer> = self.partials[..n].iter().collect(); let psize = (self.n_wg_x as u64) * (self.n_wg_y as u64) * 5 * 4;
let sizes = vec![size; n]; let mut srcs: Vec<&wgpu::Buffer> = self.partials[..n].iter().collect();
let mut sizes: Vec<u64> = vec![psize; n];
let tsize = (THUMB_W * THUMB_H * 4) as u64;
srcs.extend(self.thumb_bufs[..n].iter());
sizes.extend(vec![tsize; n]);
let raws = self.gpu.readback_many(&srcs, &sizes)?; 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_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() {
@@ -377,12 +544,30 @@ 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.next().and_then(|s| s.trim().parse::<f64>().ok()).unwrap_or(0.0); let sharp = it
let exposure = it.next().and_then(|s| s.trim().parse::<f64>().ok()).unwrap_or(0.0); .next()
let contrast = it.next().and_then(|s| s.trim().parse::<f64>().ok()).unwrap_or(0.0); .and_then(|s| s.trim().parse::<f64>().ok())
let motion = it.next().and_then(|s| s.trim().parse::<f64>().ok()).unwrap_or(0.0); .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 { frame, sharpness: sharp, exposure, contrast, motion }); out.push(FrameMetrics {
frame,
sharpness: sharp,
exposure,
contrast,
motion,
});
} }
} }
Ok(out) Ok(out)
@@ -431,7 +616,8 @@ 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> = window.iter().filter(|m| passes(cfg, m)).collect(); let survivors: Vec<&FrameMetrics> =
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())
@@ -440,12 +626,10 @@ 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 => { BadFrameAction::Warn | BadFrameAction::Keep => window
window .iter()
.iter() .max_by(|a, b| a.sharpness.partial_cmp(&b.sharpness).unwrap())
.max_by(|a, b| a.sharpness.partial_cmp(&b.sharpness).unwrap()) .map(|m| m.frame),
.map(|m| m.frame)
}
}, },
}; };
if let Some(f) = chosen { if let Some(f) = chosen {
@@ -465,7 +649,11 @@ 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 { survivors.len() as u32 } else { cfg.target_count }; let target = if cfg.target_count == 0 {
survivors.len() as u32
} else {
cfg.target_count
};
let mut accepted: Vec<u64> = Vec::new(); 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 {
@@ -490,7 +678,12 @@ 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(pattern: &str, index: u32, output_format: &str, prefix: &str) -> Result<String> { pub fn format_output_filename(
pattern: &str,
index: u32,
output_format: &str,
prefix: &str,
) -> Result<String> {
// Find the single `%[flags][width]d` token and substitute. // 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();
@@ -521,9 +714,21 @@ pub fn format_output_filename(pattern: &str, index: u32, output_format: &str, pr
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.extension().map(|e| e.to_string_lossy().to_lowercase()).as_deref() == Some("jpg") let final_name = if path
|| path.extension().map(|e| e.to_string_lossy().to_lowercase()).as_deref() == Some("jpeg") .extension()
|| path.extension().map(|e| e.to_string_lossy().to_lowercase()).as_deref() == Some("png") .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('.')) path.with_extension(target_suffix.trim_start_matches('.'))
.to_string_lossy() .to_string_lossy()
@@ -540,7 +745,14 @@ pub fn format_output_filename(pattern: &str, index: u32, output_format: &str, pr
// Frame writing // 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() { if let Some(parent) = path.parent() {
std::fs::create_dir_all(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 { 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).with_context(|| format!("cannot create {}", path.display()))?; let mut out = std::fs::File::create(path)
let mut enc = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut out, jpeg_quality as u8); .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)) 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 {
@@ -587,7 +801,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;
@@ -614,18 +828,32 @@ 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 results = analyzer.collect_batch(refs.len())?; let (metrics_b, thumbs_b) = analyzer.collect_batch(refs.len())?;
for (i, (sharp, exp, contrast)) in results.into_iter().enumerate() { for (i, (sharp, exp, contrast)) in metrics_b.into_iter().enumerate() {
let rgba = &buf[i];
let motion = if want_motion { let motion = if want_motion {
let thumb = thumbnail(rgba, w, h); // Thumbnails come back from the GPU as an `array<f32>` of luma
let m = prev_thumb.as_ref().map(|p| frame_motion(&thumb, p)).unwrap_or(0.0); // (0..255). Decode each 4-byte element to an f32, then round to
prev_thumb = Some(thumb); // u8 for the cheap CPU motion diff.
let thumb_u8: Vec<u8> = thumbs_b[i]
.chunks_exact(4)
.map(|px| f32::from_le_bytes([px[0], px[1], px[2], px[3]]).round() as u8)
.collect();
let m = prev_thumb
.as_ref()
.map(|p| frame_motion(&thumb_u8, p))
.unwrap_or(0.0);
prev_thumb = Some(thumb_u8);
m m
} else { } else {
0.0 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; frame_no += 1;
} }
if frame_no.is_multiple_of(50) || buf.len() < batch { if frame_no.is_multiple_of(50) || buf.len() < batch {
@@ -660,7 +888,12 @@ 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!("Extracting {} frames at {}x{}…", targets.len(), probe.width, probe.height), &format!(
"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)?;
@@ -670,8 +903,20 @@ 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(&cfg.output_pattern, output_index, &cfg.output_format, prefix)?; let name = format_output_filename(
save_frame(&rgba, pipe.width, pipe.height, &output_dir.join(name), cfg.jpeg_quality, &cfg.output_format)?; &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; 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);
@@ -683,7 +928,11 @@ fn extract_frames(
} }
pipe.close(); pipe.close();
if saved != targets.len() { 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) Ok(saved)
} }
@@ -693,10 +942,17 @@ fn per_video_prefixes(videos: &[PathBuf]) -> Vec<String> {
let mut used: std::collections::HashMap<String, u32> = std::collections::HashMap::new(); let mut 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.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); let n = used.entry(base.clone()).or_insert(0);
*n += 1; *n += 1;
out.push(if *n == 1 { base } else { format!("{base}_{}", *n) }); out.push(if *n == 1 {
base
} else {
format!("{base}_{}", *n)
});
} }
out out
} }
@@ -733,7 +989,15 @@ 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(&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 tag = if multi { prefix.as_str() } else { "" };
let metadata_path = output_dir.join(if multi { let metadata_path = output_dir.join(if multi {
@@ -744,18 +1008,36 @@ pub fn run_extract_stage(
let probe = video::probe(video)?; let probe = video::probe(video)?;
logging::emit( logging::emit(
&format!("Video: {}x{} {} fps, {} frames declared", &format!(
probe.width, probe.height, "Video: {}x{} {} fps, {} frames declared",
probe.fps_f64().map(|f| format!("{f:.3}")).unwrap_or_else(|| "?".into()), probe.width,
probe.nb_frames.map(|n| n.to_string()).unwrap_or_else(|| "?".into())), 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, logger,
); );
let metrics = if metadata_path.exists() && cfg.reuse_metadata { 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)? read_metadata(&metadata_path)?
} else { } 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)?; let m = analyze_video(gpu, video, &probe, cfg, should_cancel, logger)?;
write_metadata(&metadata_path, &m)?; write_metadata(&metadata_path, &m)?;
m m
@@ -766,21 +1048,46 @@ 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!("Quality filter: {n_rejected}/{} frames rejected (thresholds).", metrics.len()), &format!(
"Quality filter: {n_rejected}/{} frames rejected (thresholds).",
metrics.len()
),
logger, logger,
); );
} }
logging::emit( logging::emit(
&format!("Selected {}/{} frames (mode={:?}, chunk={}, gap={}).", &format!(
selected.len(), metrics.len(), cfg.selection, cfg.chunk_size, cfg.min_gap), "Selected {}/{} frames (mode={:?}, chunk={}, gap={}).",
selected.len(),
metrics.len(),
cfg.selection,
cfg.chunk_size,
cfg.min_gap
),
logger, logger,
); );
let saved = extract_frames(video, &probe, &selected, output_dir, cfg, tag, should_cancel, logger)?; let saved = extract_frames(
outcomes.push(ExtractOutcome { video: video.to_path_buf(), prefix: prefix.clone(), frames_written: saved }); 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("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) Ok(outcomes)
} }