1use std::io::{BufWriter, Write};
23use std::path::Path;
24use std::process::{Child, Command, Stdio};
25
26use crate::viz::Rgb8Image;
27
28#[derive(Clone, Copy, Debug, PartialEq)]
34pub struct Fps(f64);
35
36impl Fps {
37 pub fn new(fps: f64) -> Self {
40 Self(if fps.is_finite() {
41 fps.clamp(0.1, 1000.0)
42 } else {
43 30.0
44 })
45 }
46
47 pub fn get(self) -> f64 {
48 self.0
49 }
50
51 fn centiseconds(self) -> u16 {
54 ((100.0 / self.0).round() as u64).clamp(1, u16::MAX as u64) as u16
55 }
56}
57
58impl Default for Fps {
59 fn default() -> Self {
60 Self(30.0)
61 }
62}
63
64pub trait AnimationEncoder: Send {
70 fn write_frame(&mut self, image: &Rgb8Image) -> std::io::Result<()>;
72 fn finish(self: Box<Self>) -> std::io::Result<()>;
74}
75
76#[derive(Clone, Copy, Debug, PartialEq, Eq)]
78pub enum AnimationFormat {
79 Apng,
80 Gif,
81 Mp4,
82}
83
84impl AnimationFormat {
85 pub fn from_path(path: &Path) -> Option<Self> {
88 let ext = path.extension()?.to_str()?.to_ascii_lowercase();
89 Some(match ext.as_str() {
90 "apng" | "png" => Self::Apng,
91 "gif" => Self::Gif,
92 "mp4" | "m4v" | "mov" => Self::Mp4,
93 _ => return None,
94 })
95 }
96}
97
98pub fn encoder_for(
103 path: &Path,
104 frames: u32,
105 width: usize,
106 height: usize,
107 fps: Fps,
108) -> std::io::Result<Box<dyn AnimationEncoder + Send>> {
109 let format = AnimationFormat::from_path(path).ok_or_else(|| {
110 std::io::Error::new(
111 std::io::ErrorKind::InvalidInput,
112 format!(
113 "cannot infer an animation format from {}: expected .gif, .mp4 or .png/.apng",
114 path.display()
115 ),
116 )
117 })?;
118 Ok(match format {
119 AnimationFormat::Apng => Box::new(ApngEncoder::new(path, frames, width, height, fps)?)
120 as Box<dyn AnimationEncoder + Send>,
121 AnimationFormat::Gif => Box::new(GifEncoder::new(path, width, height, fps)?),
122 AnimationFormat::Mp4 => Box::new(FfmpegEncoder::new(path, width, height, fps)?),
123 })
124}
125
126pub struct ApngEncoder {
128 writer: png::Writer<BufWriter<std::fs::File>>,
129}
130
131impl ApngEncoder {
132 pub fn new(
133 path: &Path,
134 frames: u32,
135 width: usize,
136 height: usize,
137 fps: Fps,
138 ) -> std::io::Result<Self> {
139 let file = BufWriter::new(std::fs::File::create(path)?);
140 let mut encoder = png::Encoder::new(file, width as u32, height as u32);
141 encoder.set_color(png::ColorType::Rgb);
142 encoder.set_depth(png::BitDepth::Eight);
143 encoder
145 .set_animated(frames.max(1), 0)
146 .map_err(to_io_error)?;
147 encoder
149 .set_frame_delay((1000.0 / fps.get()).round() as u16, 1000)
150 .map_err(to_io_error)?;
151 let writer = encoder.write_header().map_err(to_io_error)?;
152 Ok(Self { writer })
153 }
154}
155
156impl AnimationEncoder for ApngEncoder {
157 fn write_frame(&mut self, image: &Rgb8Image) -> std::io::Result<()> {
158 self.writer
159 .write_image_data(&image.pixels)
160 .map_err(to_io_error)
161 }
162
163 fn finish(self: Box<Self>) -> std::io::Result<()> {
164 self.writer.finish().map_err(to_io_error)
165 }
166}
167
168pub struct GifEncoder {
170 encoder: gif::Encoder<BufWriter<std::fs::File>>,
171 delay: u16,
172}
173
174impl GifEncoder {
175 pub fn new(path: &Path, width: usize, height: usize, fps: Fps) -> std::io::Result<Self> {
176 let file = BufWriter::new(std::fs::File::create(path)?);
177 let mut encoder =
178 gif::Encoder::new(file, width as u16, height as u16, &[]).map_err(to_io_error)?;
179 encoder
180 .set_repeat(gif::Repeat::Infinite)
181 .map_err(to_io_error)?;
182 Ok(Self {
183 encoder,
184 delay: fps.centiseconds(),
185 })
186 }
187}
188
189impl AnimationEncoder for GifEncoder {
190 fn write_frame(&mut self, image: &Rgb8Image) -> std::io::Result<()> {
191 let mut frame =
193 gif::Frame::from_rgb(image.width as u16, image.height as u16, &image.pixels);
194 frame.delay = self.delay;
195 self.encoder.write_frame(&frame).map_err(to_io_error)
196 }
197
198 fn finish(self: Box<Self>) -> std::io::Result<()> {
199 drop(self.encoder);
201 Ok(())
202 }
203}
204
205pub struct FfmpegEncoder {
207 child: Child,
208}
209
210impl FfmpegEncoder {
211 pub fn new(path: &Path, width: usize, height: usize, fps: Fps) -> std::io::Result<Self> {
212 let child = Command::new("ffmpeg")
213 .args(["-hide_banner", "-loglevel", "error", "-y"])
214 .args(["-f", "rawvideo", "-pix_fmt", "rgb24"])
215 .args(["-s", &format!("{width}x{height}")])
216 .args(["-r", &format!("{}", fps.get())])
217 .args(["-i", "-"])
218 .args(["-vf", "pad=ceil(iw/2)*2:ceil(ih/2)*2"])
222 .args(["-c:v", "libx264", "-pix_fmt", "yuv420p"])
223 .arg(path)
224 .stdin(Stdio::piped())
225 .stdout(Stdio::null())
226 .stderr(Stdio::inherit())
227 .spawn()
228 .map_err(|error| {
229 missing_ffmpeg(
230 error,
231 "writing .mp4 needs ffmpeg on PATH",
232 "Write .gif or .apng instead to avoid the dependency.",
233 )
234 })?;
235 Ok(Self { child })
236 }
237}
238
239fn missing_ffmpeg(error: std::io::Error, what: &str, alternative: &str) -> std::io::Error {
245 if error.kind() != std::io::ErrorKind::NotFound {
246 return error;
247 }
248 std::io::Error::new(
249 std::io::ErrorKind::NotFound,
250 format!(
251 "{what} (macOS: `brew install ffmpeg`, Debian/Ubuntu: `apt install ffmpeg`, \
252 conda: `conda install -c conda-forge ffmpeg`). {alternative}"
253 )
254 .trim_end()
255 .to_owned(),
256 )
257}
258
259fn wait_for_ffmpeg(child: &mut Child) -> std::io::Result<()> {
264 let status = child.wait()?;
265 if status.success() {
266 Ok(())
267 } else {
268 Err(std::io::Error::other(format!(
269 "ffmpeg exited with {status} — its error output is above"
270 )))
271 }
272}
273
274impl AnimationEncoder for FfmpegEncoder {
275 fn write_frame(&mut self, image: &Rgb8Image) -> std::io::Result<()> {
276 let stdin = self.child.stdin.as_mut().ok_or_else(|| {
277 std::io::Error::new(std::io::ErrorKind::BrokenPipe, "ffmpeg stdin was closed")
278 })?;
279 stdin.write_all(&image.pixels)
280 }
281
282 fn finish(mut self: Box<Self>) -> std::io::Result<()> {
283 drop(self.child.stdin.take());
285 wait_for_ffmpeg(&mut self.child)
286 }
287}
288
289#[derive(Clone, Copy, Debug, PartialEq)]
294pub struct VideoInfo {
295 pub width: usize,
296 pub height: usize,
297 pub fps: f64,
298 pub frames: Option<usize>,
301}
302
303impl VideoInfo {
304 pub fn probe(path: &Path) -> std::io::Result<Self> {
306 let output = Command::new("ffprobe")
307 .args(["-v", "error", "-select_streams", "v:0"])
308 .args(["-show_entries", "stream=width,height,r_frame_rate,nb_frames"])
309 .args(["-of", "csv=p=0"])
310 .arg(path)
311 .output()
312 .map_err(|error| missing_ffmpeg(error, "reading video needs ffmpeg on PATH", ""))?;
313 if !output.status.success() {
314 return Err(std::io::Error::other(format!(
315 "ffprobe could not read {}: {}",
316 path.display(),
317 String::from_utf8_lossy(&output.stderr).trim()
318 )));
319 }
320 Self::parse(&String::from_utf8_lossy(&output.stdout), path)
321 }
322
323 fn parse(text: &str, path: &Path) -> std::io::Result<Self> {
326 let malformed = || {
327 std::io::Error::other(format!(
328 "could not read the video stream of {} — is it a video file?",
329 path.display()
330 ))
331 };
332 let line = text
333 .lines()
334 .find(|line| !line.trim().is_empty())
335 .ok_or_else(malformed)?;
336 let mut fields = line.trim().split(',');
337 let width: usize = fields
338 .next()
339 .ok_or_else(malformed)?
340 .trim()
341 .parse()
342 .map_err(|_| malformed())?;
343 let height: usize = fields
344 .next()
345 .ok_or_else(malformed)?
346 .trim()
347 .parse()
348 .map_err(|_| malformed())?;
349 let rate = fields.next().ok_or_else(malformed)?.trim();
351 let (num, den) = rate.split_once('/').unwrap_or((rate, "1"));
352 let num: f64 = num.parse().map_err(|_| malformed())?;
353 let den: f64 = den.parse().unwrap_or(1.0);
354 let fps = if den > 0.0 && num > 0.0 {
355 num / den
356 } else {
357 30.0
358 };
359 let frames = fields
362 .next()
363 .and_then(|field| field.trim().parse::<usize>().ok())
364 .filter(|&frames| frames > 0);
365 if width == 0 || height == 0 {
366 return Err(malformed());
367 }
368 Ok(Self {
369 width,
370 height,
371 fps,
372 frames,
373 })
374 }
375}
376
377pub struct FfmpegDecoder {
382 child: Child,
383 info: VideoInfo,
384 frame_bytes: usize,
385 buffer: Vec<u8>,
386 finished: bool,
387}
388
389impl FfmpegDecoder {
390 pub fn open(path: &Path, scale: Option<(usize, usize)>) -> std::io::Result<Self> {
393 let probed = VideoInfo::probe(path)?;
394 let info = match scale {
395 Some((width, height)) if width > 0 && height > 0 => VideoInfo {
396 width,
397 height,
398 fps: probed.fps,
399 frames: probed.frames,
400 },
401 _ => probed,
402 };
403 let mut command = Command::new("ffmpeg");
404 command
405 .args(["-hide_banner", "-loglevel", "error"])
406 .arg("-i")
407 .arg(path);
408 if scale.is_some() {
409 command.args(["-vf", &format!("scale={}:{}", info.width, info.height)]);
410 }
411 let child = command
412 .args(["-f", "rawvideo", "-pix_fmt", "rgb24", "-"])
413 .stdin(Stdio::null())
414 .stdout(Stdio::piped())
415 .stderr(Stdio::inherit())
416 .spawn()
417 .map_err(|error| missing_ffmpeg(error, "reading video needs ffmpeg on PATH", ""))?;
418 Ok(Self {
419 child,
420 info,
421 frame_bytes: info.width * info.height * 3,
422 buffer: vec![0; info.width * info.height * 3],
423 finished: false,
424 })
425 }
426
427 pub fn info(&self) -> VideoInfo {
428 self.info
429 }
430
431 pub fn next_frame(&mut self) -> std::io::Result<Option<Rgb8Image>> {
437 if self.finished {
438 return Ok(None);
439 }
440 let stdout = self.child.stdout.as_mut().ok_or_else(|| {
441 std::io::Error::new(std::io::ErrorKind::BrokenPipe, "ffmpeg stdout was closed")
442 })?;
443 match read_exact_or_eof(stdout, &mut self.buffer[..self.frame_bytes])? {
444 true => Ok(Some(Rgb8Image {
445 width: self.info.width,
446 height: self.info.height,
447 pixels: self.buffer[..self.frame_bytes].to_vec(),
448 })),
449 false => {
450 self.finished = true;
451 wait_for_ffmpeg(&mut self.child)?;
452 Ok(None)
453 }
454 }
455 }
456}
457
458impl Drop for FfmpegDecoder {
459 fn drop(&mut self) {
460 if !self.finished {
463 let _ = self.child.kill();
464 let _ = self.child.wait();
465 }
466 }
467}
468
469fn read_exact_or_eof(reader: &mut impl std::io::Read, buffer: &mut [u8]) -> std::io::Result<bool> {
472 let mut filled = 0;
473 while filled < buffer.len() {
474 match reader.read(&mut buffer[filled..]) {
475 Ok(0) => return Ok(false),
476 Ok(n) => filled += n,
477 Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {}
478 Err(error) => return Err(error),
479 }
480 }
481 Ok(true)
482}
483
484fn to_io_error<E: std::fmt::Display>(error: E) -> std::io::Error {
485 std::io::Error::other(error.to_string())
486}
487
488#[cfg(test)]
489mod tests {
490 use super::*;
491 use crate::viz::Rgb8Image;
492
493 fn frame(width: usize, height: usize, shade: u8) -> Rgb8Image {
494 Rgb8Image {
495 width,
496 height,
497 pixels: vec![shade; width * height * 3],
498 }
499 }
500
501 fn temp_path(name: &str) -> std::path::PathBuf {
502 let mut path = std::env::temp_dir();
503 path.push(format!(
504 "eventcv-video-test-{}-{}",
505 std::process::id(),
506 name
507 ));
508 path
509 }
510
511 #[test]
512 fn format_is_read_from_the_extension() {
513 let cases = [
514 ("a.gif", Some(AnimationFormat::Gif)),
515 ("a.GIF", Some(AnimationFormat::Gif)),
516 ("a.png", Some(AnimationFormat::Apng)),
517 ("a.apng", Some(AnimationFormat::Apng)),
518 ("a.mp4", Some(AnimationFormat::Mp4)),
519 ("a.mov", Some(AnimationFormat::Mp4)),
520 ("a.txt", None),
521 ("a", None),
522 ];
523 for (name, expected) in cases {
524 assert_eq!(
525 AnimationFormat::from_path(Path::new(name)),
526 expected,
527 "{name}"
528 );
529 }
530 }
531
532 #[test]
533 fn fps_clamps_and_converts() {
534 assert_eq!(Fps::new(f64::NAN).get(), 30.0);
535 assert_eq!(Fps::new(0.0).get(), 0.1);
536 assert_eq!(Fps::new(1e9).get(), 1000.0);
537 assert_eq!(Fps::new(100.0).centiseconds(), 1); assert_eq!(Fps::new(50.0).centiseconds(), 2);
539 assert_eq!(Fps::new(10.0).centiseconds(), 10);
540 }
541
542 #[test]
543 fn apng_writes_a_multi_frame_file() {
544 let path = temp_path("apng.png");
545 let mut encoder: Box<dyn AnimationEncoder> =
546 Box::new(ApngEncoder::new(&path, 3, 4, 4, Fps::new(10.0)).unwrap());
547 for shade in [0u8, 128, 255] {
548 encoder.write_frame(&frame(4, 4, shade)).unwrap();
549 }
550 encoder.finish().unwrap();
551
552 let bytes = std::fs::read(&path).unwrap();
553 assert_eq!(&bytes[..8], b"\x89PNG\r\n\x1a\n");
554 assert!(bytes.windows(4).any(|w| w == b"acTL"));
557 assert!(bytes.windows(4).any(|w| w == b"fcTL"));
558 std::fs::remove_file(&path).ok();
559 }
560
561 #[test]
562 fn gif_writes_a_multi_frame_file() {
563 let path = temp_path("gif.gif");
564 let mut encoder: Box<dyn AnimationEncoder> =
565 Box::new(GifEncoder::new(&path, 4, 4, Fps::new(10.0)).unwrap());
566 for shade in [0u8, 128, 255] {
567 encoder.write_frame(&frame(4, 4, shade)).unwrap();
568 }
569 encoder.finish().unwrap();
570
571 let bytes = std::fs::read(&path).unwrap();
572 assert_eq!(&bytes[..6], b"GIF89a");
573 assert_eq!(bytes.last(), Some(&0x3B)); std::fs::remove_file(&path).ok();
575 }
576
577 #[test]
578 fn video_info_parses_ffprobe_csv() {
579 let path = Path::new("clip.mp4");
580 let info = VideoInfo::parse("64,48,30/1\n", path).unwrap();
582 assert_eq!((info.width, info.height), (64, 48));
583 assert!((info.fps - 30.0).abs() < 1e-9);
584 let ntsc = VideoInfo::parse("1920,1080,30000/1001", path).unwrap();
586 assert!((ntsc.fps - 29.97).abs() < 0.01);
587 assert!((VideoInfo::parse("8,8,25", path).unwrap().fps - 25.0).abs() < 1e-9);
589 }
590
591 #[test]
592 fn video_info_rejects_nonsense() {
593 let path = Path::new("notes.txt");
594 for text in ["", "\n", "not,a,video", "0,0,30/1"] {
595 assert!(VideoInfo::parse(text, path).is_err(), "{text:?}");
596 }
597 }
598
599 #[test]
600 fn decoder_reads_back_every_frame_it_was_given() {
601 if Command::new("ffmpeg").arg("-version").output().is_err() {
604 return; }
606 let path = temp_path("roundtrip.mp4");
607 let mut encoder: Box<dyn AnimationEncoder + Send> =
608 Box::new(FfmpegEncoder::new(&path, 32, 24, Fps::new(10.0)).unwrap());
609 for shade in [0u8, 60, 120, 180, 240] {
610 encoder.write_frame(&frame(32, 24, shade)).unwrap();
611 }
612 encoder.finish().unwrap();
613
614 let mut decoder = FfmpegDecoder::open(&path, None).unwrap();
615 assert_eq!((decoder.info().width, decoder.info().height), (32, 24));
616 let mut decoded = 0;
617 while let Some(image) = decoder.next_frame().unwrap() {
618 assert_eq!(image.pixels.len(), 32 * 24 * 3);
619 decoded += 1;
620 }
621 assert_eq!(decoded, 5);
622 assert!(decoder.next_frame().unwrap().is_none());
624 std::fs::remove_file(&path).ok();
625 }
626
627 #[test]
628 fn decoder_can_scale_on_the_way_out() {
629 if Command::new("ffmpeg").arg("-version").output().is_err() {
630 return;
631 }
632 let path = temp_path("scaled.mp4");
633 let mut encoder: Box<dyn AnimationEncoder + Send> =
634 Box::new(FfmpegEncoder::new(&path, 64, 64, Fps::new(10.0)).unwrap());
635 encoder.write_frame(&frame(64, 64, 128)).unwrap();
636 encoder.finish().unwrap();
637
638 let mut decoder = FfmpegDecoder::open(&path, Some((16, 16))).unwrap();
639 let image = decoder.next_frame().unwrap().expect("one frame");
640 assert_eq!((image.width, image.height), (16, 16));
641 assert_eq!(image.pixels.len(), 16 * 16 * 3);
642 std::fs::remove_file(&path).ok();
643 }
644
645 #[test]
646 fn read_exact_or_eof_reports_a_clean_end() {
647 let mut full = [0u8; 4];
648 assert!(read_exact_or_eof(&mut &b"abcd"[..], &mut full).unwrap());
649 assert_eq!(&full, b"abcd");
650 assert!(!read_exact_or_eof(&mut &b""[..], &mut full).unwrap());
652 assert!(!read_exact_or_eof(&mut &b"ab"[..], &mut full).unwrap());
654 }
655
656 #[test]
657 fn unknown_extension_is_rejected_with_a_useful_message() {
658 let error = encoder_for(Path::new("out.avi"), 1, 4, 4, Fps::default())
659 .err()
660 .expect("an unknown extension must not open an encoder");
661 assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
662 assert!(error.to_string().contains(".gif"));
663 }
664
665 #[test]
666 fn missing_ffmpeg_names_the_fix() {
667 if Command::new("ffmpeg").arg("-version").output().is_err() {
670 let error = FfmpegEncoder::new(&temp_path("x.mp4"), 4, 4, Fps::default())
671 .err()
672 .expect("spawning ffmpeg must fail when it is not installed");
673 assert_eq!(error.kind(), std::io::ErrorKind::NotFound);
674 assert!(error.to_string().contains("ffmpeg"));
675 assert!(error.to_string().contains(".gif"));
676 }
677 }
678}