FEMM/crates/femm-app/src/export.rs

617 lines
22 KiB
Rust

//! offline RGBA frame renderer piping to an ffmpeg VP9 encoder.
use std::collections::HashMap;
use std::io::{Read, Write};
use std::path::PathBuf;
use std::process::{Command, Stdio};
use std::sync::{Arc, Mutex};
use std::thread;
use femm_doc_mag::FemmDoc;
use femm_doc_mag::ans::MagSolution;
use tiny_skia::{Color, FillRule, Paint, PathBuilder, Pixmap, Stroke, Transform};
use crate::kinematic::{self, Track};
const BAND_COUNT: usize = 20;
const FLUX_LINE_COUNT: usize = 19;
const PADDING_PX: f64 = 48.0;
pub struct ExportInput {
pub dest: PathBuf,
pub base_doc: FemmDoc,
pub tracks: Vec<Track>,
pub dt: f64,
pub buffer_size: usize,
pub fps: f64,
pub lut: HashMap<i64, MagSolution>,
pub width: u32,
pub height: u32,
pub contour: bool,
pub crf: u32,
pub probes: Vec<crate::probe::Probe>,
pub probe_samples: Vec<HashMap<i64, (f64, f64)>>,
pub progress: Arc<Mutex<ExportProgress>>,
}
#[derive(Debug, Clone, Default)]
pub struct ExportProgress {
pub phase: ExportPhase,
pub frames_done: usize,
pub frames_total: usize,
pub bytes_sent: usize,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub enum ExportPhase {
#[default]
Idle,
Rendering,
Finalizing,
Done,
}
#[derive(Debug, Clone)]
pub struct ExportReport {
pub written: usize,
pub dest: PathBuf,
}
/// renders each buffered frame to a pixmap and pipes the raw RGBA bytes to a libvpx-vp9 ffmpeg subprocess.
pub fn export_webm(input: ExportInput) -> Result<ExportReport, String> {
if input.buffer_size == 0 {
return Err(String::from("buffer is empty - set a buffer size and run the sim first"));
}
for f in 0..(input.buffer_size as i64) {
if !input.lut.contains_key(&f) {
return Err(format!("frame {f} of {} not cached - fill the buffer before exporting", input.buffer_size));
}
}
let size_arg = format!("{}x{}", input.width, input.height);
let fps_arg = format!("{:.6}", input.fps.max(1.0));
let dest_str = input.dest.to_string_lossy().to_string();
let ffmpeg = locate_ffmpeg().ok_or_else(|| String::from(
"ffmpeg binary not found. install via `brew install ffmpeg`, or set FFMPEG_PATH to its absolute path. searched: $FFMPEG_PATH, $PATH, /opt/homebrew/bin/ffmpeg, /usr/local/bin/ffmpeg, /opt/local/bin/ffmpeg, /usr/bin/ffmpeg"
))?;
let crf_arg = input.crf.clamp(0, 63).to_string();
let mut ff = Command::new(&ffmpeg)
.args([
"-y",
"-loglevel", "info",
"-f", "rawvideo",
"-pixel_format", "rgba",
"-video_size", size_arg.as_str(),
"-framerate", fps_arg.as_str(),
"-i", "-",
"-c:v", "libvpx-vp9",
"-pix_fmt", "yuv420p",
"-crf", crf_arg.as_str(),
"-b:v", "0",
"-r", fps_arg.as_str(),
"-fps_mode", "cfr",
"-vsync", "cfr",
"-deadline", "good",
"-cpu-used", "2",
"-row-mt", "1",
"-tile-columns", "2",
"-threads", "8",
dest_str.as_str(),
])
.stdin(Stdio::piped())
.stdout(Stdio::null())
.stderr(Stdio::piped())
.spawn()
.map_err(|e| format!("ffmpeg spawn failed: {e}"))?;
let mut stdin = ff.stdin.take().ok_or_else(|| String::from("ffmpeg stdin pipe missing"))?;
let stderr_buf = Arc::new(Mutex::new(Vec::<u8>::new()));
let stderr_pipe = ff.stderr.take().ok_or_else(|| String::from("ffmpeg stderr pipe missing"))?;
let stderr_thread = {
let buf = Arc::clone(&stderr_buf);
thread::spawn(move || {
let mut reader = stderr_pipe;
let mut chunk = [0u8; 4096];
loop {
match reader.read(&mut chunk) {
Ok(0) | Err(_) => break,
Ok(n) => buf.lock().unwrap().extend_from_slice(&chunk[..n]),
}
}
})
};
let (xmin, xmax, ymin, ymax) = world_bounds(&input.base_doc);
let expected_bytes_per_frame = (input.width as usize) * (input.height as usize) * 4;
{
let mut p = input.progress.lock().unwrap();
p.phase = ExportPhase::Rendering;
p.frames_done = 0;
p.frames_total = input.buffer_size;
p.bytes_sent = 0;
}
let mut written = 0usize;
let mut bytes_sent = 0usize;
for f in 0..(input.buffer_size as i64) {
let sol = input.lut.get(&f).expect("frame presence verified above");
let mut doc = input.base_doc.clone();
let t = (f as f64) * input.dt;
kinematic::apply_tracks(&mut doc, &input.base_doc, &input.tracks, t);
let pixmap = render_frame(
&doc, sol, input.width, input.height, input.contour,
xmin, xmax, ymin, ymax,
&input.probes, &input.probe_samples,
f, input.buffer_size,
);
let frame_bytes = pixmap.data();
if frame_bytes.len() != expected_bytes_per_frame {
drop(stdin);
let _ = ff.kill();
let _ = stderr_thread.join();
return Err(format!(
"frame {f} produced {} bytes, expected {} (width={}, height={})",
frame_bytes.len(), expected_bytes_per_frame, input.width, input.height,
));
}
if let Err(e) = stdin.write_all(frame_bytes) {
drop(stdin);
let _ = ff.wait();
let _ = stderr_thread.join();
let log = String::from_utf8_lossy(&stderr_buf.lock().unwrap()).to_string();
return Err(format!(
"write frame {f} to ffmpeg failed after {bytes_sent} bytes ({written} frames): {e}\n\n--- ffmpeg log ---\n{log}",
));
}
bytes_sent += frame_bytes.len();
written += 1;
{
let mut p = input.progress.lock().unwrap();
p.frames_done = written;
p.bytes_sent = bytes_sent;
}
}
drop(stdin);
{
let mut p = input.progress.lock().unwrap();
p.phase = ExportPhase::Finalizing;
}
let status = ff.wait().map_err(|e| format!("ffmpeg wait failed: {e}"))?;
let _ = stderr_thread.join();
let log = String::from_utf8_lossy(&stderr_buf.lock().unwrap()).to_string();
let dest_size = std::fs::metadata(&input.dest).map(|m| m.len()).unwrap_or(0);
if !status.success() {
return Err(format!(
"ffmpeg exited with {status}. wrote {written} frames ({bytes_sent} bytes) to stdin, output {dest_size} bytes.\n\n--- ffmpeg log ---\n{log}",
));
}
if dest_size == 0 {
return Err(format!(
"ffmpeg succeeded but the output is 0 bytes. wrote {written} frames ({bytes_sent} bytes) to stdin at {input_w}x{input_h} rgba, {fps} fps.\n\n--- ffmpeg log ---\n{log}",
input_w = input.width, input_h = input.height, fps = fps_arg,
));
}
{
let mut p = input.progress.lock().unwrap();
p.phase = ExportPhase::Done;
}
Ok(ExportReport { written, dest: input.dest })
}
/// locates an ffmpeg binary via $FFMPEG_PATH, $PATH, and the standard Homebrew, MacPorts, and system locations.
fn locate_ffmpeg() -> Option<std::path::PathBuf> {
if let Ok(p) = std::env::var("FFMPEG_PATH") {
let pb = std::path::PathBuf::from(p);
if pb.is_file() { return Some(pb); }
}
if let Ok(path) = std::env::var("PATH") {
for dir in path.split(':') {
let pb = std::path::PathBuf::from(dir).join("ffmpeg");
if pb.is_file() { return Some(pb); }
}
}
for candidate in [
"/opt/homebrew/bin/ffmpeg",
"/usr/local/bin/ffmpeg",
"/opt/local/bin/ffmpeg",
"/usr/bin/ffmpeg",
] {
let pb = std::path::PathBuf::from(candidate);
if pb.is_file() { return Some(pb); }
}
None
}
fn world_bounds(doc: &FemmDoc) -> (f64, f64, f64, f64) {
let mut xmin = f64::INFINITY;
let mut xmax = f64::NEG_INFINITY;
let mut ymin = f64::INFINITY;
let mut ymax = f64::NEG_INFINITY;
let mut had_point = false;
for n in &doc.nodes {
if n.x < xmin { xmin = n.x; }
if n.x > xmax { xmax = n.x; }
if n.y < ymin { ymin = n.y; }
if n.y > ymax { ymax = n.y; }
had_point = true;
}
for l in &doc.block_labels {
if l.x < xmin { xmin = l.x; }
if l.x > xmax { xmax = l.x; }
if l.y < ymin { ymin = l.y; }
if l.y > ymax { ymax = l.y; }
had_point = true;
}
if !had_point { return (-1.0, 1.0, -1.0, 1.0); }
if (xmax - xmin).abs() < 1e-9 { xmin -= 0.5; xmax += 0.5; }
if (ymax - ymin).abs() < 1e-9 { ymin -= 0.5; ymax += 0.5; }
(xmin, xmax, ymin, ymax)
}
struct View {
scale: f64,
offset_x: f64,
offset_y: f64,
y_max: f64,
}
impl View {
fn fit(width: u32, height: u32, xmin: f64, xmax: f64, ymin: f64, ymax: f64) -> Self {
let avail_w = (width as f64 - 2.0 * PADDING_PX).max(1.0);
let avail_h = (height as f64 - 2.0 * PADDING_PX).max(1.0);
let dx = (xmax - xmin).max(1e-9);
let dy = (ymax - ymin).max(1e-9);
let scale = (avail_w / dx).min(avail_h / dy);
let drawn_w = dx * scale;
let drawn_h = dy * scale;
let offset_x = (width as f64 - drawn_w) / 2.0 - xmin * scale;
let offset_y = (height as f64 - drawn_h) / 2.0;
let _ = ymin;
Self { scale, offset_x, offset_y, y_max: ymax }
}
fn map(&self, x: f64, y: f64) -> (f32, f32) {
let px = x * self.scale + self.offset_x;
let py = (self.y_max - y) * self.scale + self.offset_y;
(px as f32, py as f32)
}
}
fn render_frame(
doc: &FemmDoc,
sol: &MagSolution,
width: u32,
height: u32,
contour: bool,
xmin: f64, xmax: f64, ymin: f64, ymax: f64,
probes: &[crate::probe::Probe],
probe_samples: &[HashMap<i64, (f64, f64)>],
current_frame: i64,
buffer_size: usize,
) -> Pixmap {
let mut pixmap = Pixmap::new(width, height).expect("allocate pixmap");
pixmap.fill(Color::from_rgba8(20, 20, 24, 255));
let plot_strip_h = if probes.is_empty() { 0 } else { (height as f32 * 0.33) as u32 };
let canvas_h = height - plot_strip_h;
let view = View::fit(width, canvas_h, xmin, xmax, ymin, ymax);
let (lo, hi) = sol.b_magnitude_range();
let span = if hi > lo { hi - lo } else { 1.0 };
let mut paint = Paint::default();
paint.anti_alias = true;
for i in 0..sol.mesh_elements.len() {
let el = &sol.mesh_elements[i];
let (Some(a), Some(b), Some(c)) = (
sol.mesh_nodes.get(el.n[0] as usize),
sol.mesh_nodes.get(el.n[1] as usize),
sol.mesh_nodes.get(el.n[2] as usize),
) else { continue };
let (pa_x, pa_y) = view.map(a.x, a.y);
let (pb_x, pb_y) = view.map(b.x, b.y);
let (pc_x, pc_y) = view.map(c.x, c.y);
let t = ((sol.b_magnitude(i) - lo) / span).clamp(0.0, 1.0);
let t_mapped = if contour {
let band = (t * BAND_COUNT as f64).floor().min(BAND_COUNT as f64 - 1.0);
(band + 0.5) / BAND_COUNT as f64
} else {
t
};
paint.set_color(jet_color(t_mapped as f32));
let mut pb = PathBuilder::new();
pb.move_to(pa_x, pa_y);
pb.line_to(pb_x, pb_y);
pb.line_to(pc_x, pc_y);
pb.close();
if let Some(path) = pb.finish() {
pixmap.fill_path(&path, &paint, FillRule::Winding, Transform::identity(), None);
}
}
if contour {
let (a_lo, a_hi) = sol.a_real_range();
if a_hi > a_lo {
let mut levels = Vec::with_capacity(FLUX_LINE_COUNT);
let denom = (FLUX_LINE_COUNT + 1) as f64;
for k in 1..=FLUX_LINE_COUNT {
levels.push(a_lo + (a_hi - a_lo) * (k as f64) / denom);
}
paint.set_color(Color::from_rgba8(20, 20, 24, 220));
let stroke = Stroke { width: 1.0, ..Default::default() };
for (p0, p1) in sol.flux_lines(&levels) {
let (ax, ay) = view.map(p0.0, p0.1);
let (bx, by) = view.map(p1.0, p1.1);
let mut pb = PathBuilder::new();
pb.move_to(ax, ay);
pb.line_to(bx, by);
if let Some(path) = pb.finish() {
pixmap.stroke_path(&path, &paint, &stroke, Transform::identity(), None);
}
}
}
}
paint.set_color(Color::from_rgba8(230, 230, 235, 255));
let stroke = Stroke { width: 1.4, ..Default::default() };
for s in &doc.segments {
if let (Some(p0), Some(p1)) = (doc.nodes.get(s.n0 as usize), doc.nodes.get(s.n1 as usize)) {
let (ax, ay) = view.map(p0.x, p0.y);
let (bx, by) = view.map(p1.x, p1.y);
let mut pb = PathBuilder::new();
pb.move_to(ax, ay);
pb.line_to(bx, by);
if let Some(path) = pb.finish() {
pixmap.stroke_path(&path, &paint, &stroke, Transform::identity(), None);
}
}
}
for a in &doc.arcs {
let (Some(p0), Some(p1)) = (doc.nodes.get(a.n0 as usize), doc.nodes.get(a.n1 as usize)) else { continue };
let polyline = arc_polyline(p0.x, p0.y, p1.x, p1.y, a.arc_length, 24);
let mut pb = PathBuilder::new();
for (i, (x, y)) in polyline.iter().enumerate() {
let (px, py) = view.map(*x, *y);
if i == 0 { pb.move_to(px, py); } else { pb.line_to(px, py); }
}
if let Some(path) = pb.finish() {
pixmap.stroke_path(&path, &paint, &stroke, Transform::identity(), None);
}
}
for probe in probes {
let (px, py) = view.map(probe.x, probe.y);
let theta = probe.angle_deg.to_radians() as f32;
let pri_len = 22.0_f32;
let per_len = 12.0_f32;
let pri_ex = px + theta.cos() * pri_len;
let pri_ey = py - theta.sin() * pri_len;
let per_ex = px + (theta - std::f32::consts::FRAC_PI_2).cos() * per_len;
let per_ey = py - (theta - std::f32::consts::FRAC_PI_2).sin() * per_len;
paint.set_color(Color::from_rgba8(probe.color[0], probe.color[1], probe.color[2], 255));
let pri_stroke = Stroke { width: 2.2, ..Default::default() };
let mut pb = PathBuilder::new();
pb.move_to(px, py);
pb.line_to(pri_ex, pri_ey);
if let Some(path) = pb.finish() {
pixmap.stroke_path(&path, &paint, &pri_stroke, Transform::identity(), None);
}
paint.set_color(Color::from_rgba8(probe.color[0], probe.color[1], probe.color[2], 140));
let per_stroke = Stroke { width: 1.2, ..Default::default() };
let mut pb = PathBuilder::new();
pb.move_to(px, py);
pb.line_to(per_ex, per_ey);
if let Some(path) = pb.finish() {
pixmap.stroke_path(&path, &paint, &per_stroke, Transform::identity(), None);
}
paint.set_color(Color::from_rgba8(probe.color[0], probe.color[1], probe.color[2], 255));
let mut pb = PathBuilder::new();
pb.push_circle(px, py, 6.0);
if let Some(path) = pb.finish() {
pixmap.fill_path(&path, &paint, FillRule::Winding, Transform::identity(), None);
}
paint.set_color(Color::from_rgba8(255, 255, 255, 255));
let ring = Stroke { width: 1.6, ..Default::default() };
let mut pb = PathBuilder::new();
pb.push_circle(px, py, 6.0);
if let Some(path) = pb.finish() {
pixmap.stroke_path(&path, &paint, &ring, Transform::identity(), None);
}
}
if plot_strip_h > 0 && !probes.is_empty() {
draw_plot_strip(
&mut pixmap,
0, canvas_h, width, plot_strip_h,
probes, probe_samples,
current_frame, buffer_size,
);
}
pixmap
}
/// draws one plot pane per probe across a horizontal strip at (x0, y0, w, h), each pane color-matched to its probe and marked with a vertical bar at the current frame.
fn draw_plot_strip(
pixmap: &mut Pixmap,
x0: u32, y0: u32, w: u32, h: u32,
probes: &[crate::probe::Probe],
probe_samples: &[HashMap<i64, (f64, f64)>],
current_frame: i64,
buffer_size: usize,
) {
let n = probes.len().max(1);
let pane_w = (w as f32) / (n as f32);
let pane_x_top_pad = 4.0;
let pane_y_top_pad = 18.0;
let pane_bot_pad = 6.0;
let pane_lr_pad = 4.0;
let mut paint = Paint::default();
paint.anti_alias = true;
paint.set_color(Color::from_rgba8(12, 14, 20, 255));
let strip = tiny_skia::Rect::from_xywh(x0 as f32, y0 as f32, w as f32, h as f32);
if let Some(rect) = strip {
pixmap.fill_rect(rect, &paint, Transform::identity(), None);
}
for (i, probe) in probes.iter().enumerate() {
let px0 = x0 as f32 + (i as f32) * pane_w + pane_lr_pad;
let py0 = y0 as f32 + pane_x_top_pad;
let pw = pane_w - 2.0 * pane_lr_pad;
let ph = (h as f32) - pane_x_top_pad - pane_bot_pad;
if pw <= 0.0 || ph <= 0.0 { continue; }
paint.set_color(Color::from_rgba8(18, 20, 26, 255));
if let Some(r) = tiny_skia::Rect::from_xywh(px0, py0, pw, ph) {
pixmap.fill_rect(r, &paint, Transform::identity(), None);
}
let plot_x = px0 + 4.0;
let plot_y = py0 + pane_y_top_pad;
let plot_w = pw - 8.0;
let plot_h = ph - pane_y_top_pad - 6.0;
if plot_w <= 0.0 || plot_h <= 0.0 { continue; }
let probe_color = Color::from_rgba8(probe.color[0], probe.color[1], probe.color[2], 255);
paint.set_color(probe_color);
let samples = probe_samples.get(i);
let mut points: Vec<(i64, f64)> = Vec::new();
if let Some(s) = samples {
for (&fi, &(bx, by)) in s {
points.push((fi, probe.mode.extract(bx, by, probe.angle_deg.to_radians())));
}
}
if points.is_empty() {
continue;
}
let mut ymin = f64::INFINITY;
let mut ymax = f64::NEG_INFINITY;
for &(_, v) in &points {
if v < ymin { ymin = v; }
if v > ymax { ymax = v; }
}
if !ymin.is_finite() || !ymax.is_finite() { ymin = -1.0; ymax = 1.0; }
if (ymax - ymin).abs() < 1e-12 { ymin -= 1.0; ymax += 1.0; }
points.sort_by_key(|p| p.0);
let buf = buffer_size.max(1) as f64;
let x_for = |fi: i64| -> f32 {
let t = (fi.rem_euclid(buf as i64) as f64) / buf;
plot_x + (t as f32) * plot_w
};
let y_for = |v: f64| -> f32 {
let t = (v - ymin) / (ymax - ymin);
plot_y + plot_h - (t as f32) * plot_h
};
let axis_stroke = Stroke { width: 1.0, ..Default::default() };
paint.set_color(Color::from_rgba8(60, 65, 75, 255));
let mut pb = PathBuilder::new();
pb.move_to(plot_x, plot_y + plot_h);
pb.line_to(plot_x + plot_w, plot_y + plot_h);
if let Some(path) = pb.finish() {
pixmap.stroke_path(&path, &paint, &axis_stroke, Transform::identity(), None);
}
let mid_y = y_for((ymin + ymax) * 0.5);
paint.set_color(Color::from_rgba8(50, 55, 65, 255));
let mut pb = PathBuilder::new();
pb.move_to(plot_x, mid_y);
pb.line_to(plot_x + plot_w, mid_y);
if let Some(path) = pb.finish() {
pixmap.stroke_path(&path, &paint, &axis_stroke, Transform::identity(), None);
}
paint.set_color(probe_color);
let line_stroke = Stroke { width: 1.4, ..Default::default() };
let mut pb = PathBuilder::new();
let mut started = false;
for (fi, v) in &points {
let x = x_for(*fi);
let y = y_for(*v);
if !started { pb.move_to(x, y); started = true; }
else { pb.line_to(x, y); }
}
if started {
if let Some(path) = pb.finish() {
pixmap.stroke_path(&path, &paint, &line_stroke, Transform::identity(), None);
}
}
let cx = x_for(current_frame);
paint.set_color(Color::from_rgba8(255, 255, 255, 200));
let cur_stroke = Stroke { width: 1.2, ..Default::default() };
let mut pb = PathBuilder::new();
pb.move_to(cx, plot_y);
pb.line_to(cx, plot_y + plot_h);
if let Some(path) = pb.finish() {
pixmap.stroke_path(&path, &paint, &cur_stroke, Transform::identity(), None);
}
let label_dot_r = 4.0;
paint.set_color(probe_color);
let mut pb = PathBuilder::new();
pb.push_circle(px0 + 8.0, py0 + 9.0, label_dot_r);
if let Some(path) = pb.finish() {
pixmap.fill_path(&path, &paint, FillRule::Winding, Transform::identity(), None);
}
}
}
/// returns a polyline approximating an arc between two endpoints across a CCW degree sweep, capped at max_steps subdivisions.
fn arc_polyline(x0: f64, y0: f64, x1: f64, y1: f64, sweep_deg: f64, max_steps: usize) -> Vec<(f64, f64)> {
let theta = sweep_deg.to_radians();
if theta.abs() < 1e-9 { return vec![(x0, y0), (x1, y1)]; }
let dx = x1 - x0;
let dy = y1 - y0;
let chord = (dx * dx + dy * dy).sqrt();
if chord < 1e-12 { return vec![(x0, y0)]; }
let radius = chord / (2.0 * (theta / 2.0).sin().abs());
let mid_x = (x0 + x1) / 2.0;
let mid_y = (y0 + y1) / 2.0;
let perp_x = -dy / chord;
let perp_y = dx / chord;
let h = radius * (theta / 2.0).cos();
let sign = if theta > 0.0 { 1.0 } else { -1.0 };
let cx = mid_x + sign * perp_x * h;
let cy = mid_y + sign * perp_y * h;
let a0 = (y0 - cy).atan2(x0 - cx);
let a1 = (y1 - cy).atan2(x1 - cx);
let mut delta = a1 - a0;
if theta > 0.0 {
while delta < 0.0 { delta += std::f64::consts::TAU; }
} else {
while delta > 0.0 { delta -= std::f64::consts::TAU; }
}
let steps = ((sweep_deg.abs() / 4.0).ceil() as usize).clamp(2, max_steps);
let mut out = Vec::with_capacity(steps + 1);
for i in 0..=steps {
let a = a0 + delta * (i as f64) / (steps as f64);
out.push((cx + radius * a.cos(), cy + radius * a.sin()));
}
out
}
fn jet_color(t: f32) -> Color {
let t = t.clamp(0.0, 1.0);
let (r, g, b) = if t < 0.25 {
let s = t * 4.0;
(0.0, s, 1.0)
} else if t < 0.5 {
let s = (t - 0.25) * 4.0;
(0.0, 1.0, 1.0 - s)
} else if t < 0.75 {
let s = (t - 0.5) * 4.0;
(s, 1.0, 0.0)
} else {
let s = (t - 0.75) * 4.0;
(1.0, 1.0 - s, 0.0)
};
Color::from_rgba8((r * 255.0) as u8, (g * 255.0) as u8, (b * 255.0) as u8, 255)
}