use std::io::{BufWriter, Write};
use std::path::Path;
use std::process::{Child, Command, Stdio};
use crate::viz::Rgb8Image;
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Fps(f64);
impl Fps {
pub fn new(fps: f64) -> Self {
Self(if fps.is_finite() {
fps.clamp(0.1, 1000.0)
} else {
30.0
})
}
pub fn get(self) -> f64 {
self.0
}
fn centiseconds(self) -> u16 {
((100.0 / self.0).round() as u64).clamp(1, u16::MAX as u64) as u16
}
}
impl Default for Fps {
fn default() -> Self {
Self(30.0)
}
}
pub trait AnimationEncoder: Send {
fn write_frame(&mut self, image: &Rgb8Image) -> std::io::Result<()>;
fn finish(self: Box<Self>) -> std::io::Result<()>;
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AnimationFormat {
Apng,
Gif,
Mp4,
}
impl AnimationFormat {
pub fn from_path(path: &Path) -> Option<Self> {
let ext = path.extension()?.to_str()?.to_ascii_lowercase();
Some(match ext.as_str() {
"apng" | "png" => Self::Apng,
"gif" => Self::Gif,
"mp4" | "m4v" | "mov" => Self::Mp4,
_ => return None,
})
}
}
pub fn encoder_for(
path: &Path,
frames: u32,
width: usize,
height: usize,
fps: Fps,
) -> std::io::Result<Box<dyn AnimationEncoder + Send>> {
let format = AnimationFormat::from_path(path).ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"cannot infer an animation format from {}: expected .gif, .mp4 or .png/.apng",
path.display()
),
)
})?;
Ok(match format {
AnimationFormat::Apng => Box::new(ApngEncoder::new(path, frames, width, height, fps)?)
as Box<dyn AnimationEncoder + Send>,
AnimationFormat::Gif => Box::new(GifEncoder::new(path, width, height, fps)?),
AnimationFormat::Mp4 => Box::new(FfmpegEncoder::new(path, width, height, fps)?),
})
}
pub struct ApngEncoder {
writer: png::Writer<BufWriter<std::fs::File>>,
}
impl ApngEncoder {
pub fn new(
path: &Path,
frames: u32,
width: usize,
height: usize,
fps: Fps,
) -> std::io::Result<Self> {
let file = BufWriter::new(std::fs::File::create(path)?);
let mut encoder = png::Encoder::new(file, width as u32, height as u32);
encoder.set_color(png::ColorType::Rgb);
encoder.set_depth(png::BitDepth::Eight);
encoder
.set_animated(frames.max(1), 0)
.map_err(to_io_error)?;
encoder
.set_frame_delay((1000.0 / fps.get()).round() as u16, 1000)
.map_err(to_io_error)?;
let writer = encoder.write_header().map_err(to_io_error)?;
Ok(Self { writer })
}
}
impl AnimationEncoder for ApngEncoder {
fn write_frame(&mut self, image: &Rgb8Image) -> std::io::Result<()> {
self.writer
.write_image_data(&image.pixels)
.map_err(to_io_error)
}
fn finish(self: Box<Self>) -> std::io::Result<()> {
self.writer.finish().map_err(to_io_error)
}
}
pub struct GifEncoder {
encoder: gif::Encoder<BufWriter<std::fs::File>>,
delay: u16,
}
impl GifEncoder {
pub fn new(path: &Path, width: usize, height: usize, fps: Fps) -> std::io::Result<Self> {
let file = BufWriter::new(std::fs::File::create(path)?);
let mut encoder =
gif::Encoder::new(file, width as u16, height as u16, &[]).map_err(to_io_error)?;
encoder
.set_repeat(gif::Repeat::Infinite)
.map_err(to_io_error)?;
Ok(Self {
encoder,
delay: fps.centiseconds(),
})
}
}
impl AnimationEncoder for GifEncoder {
fn write_frame(&mut self, image: &Rgb8Image) -> std::io::Result<()> {
let mut frame =
gif::Frame::from_rgb(image.width as u16, image.height as u16, &image.pixels);
frame.delay = self.delay;
self.encoder.write_frame(&frame).map_err(to_io_error)
}
fn finish(self: Box<Self>) -> std::io::Result<()> {
drop(self.encoder);
Ok(())
}
}
pub struct FfmpegEncoder {
child: Child,
}
impl FfmpegEncoder {
pub fn new(path: &Path, width: usize, height: usize, fps: Fps) -> std::io::Result<Self> {
let child = Command::new("ffmpeg")
.args(["-hide_banner", "-loglevel", "error", "-y"])
.args(["-f", "rawvideo", "-pix_fmt", "rgb24"])
.args(["-s", &format!("{width}x{height}")])
.args(["-r", &format!("{}", fps.get())])
.args(["-i", "-"])
.args(["-vf", "pad=ceil(iw/2)*2:ceil(ih/2)*2"])
.args(["-c:v", "libx264", "-pix_fmt", "yuv420p"])
.arg(path)
.stdin(Stdio::piped())
.stdout(Stdio::null())
.stderr(Stdio::inherit())
.spawn()
.map_err(|error| {
missing_ffmpeg(
error,
"writing .mp4 needs ffmpeg on PATH",
"Write .gif or .apng instead to avoid the dependency.",
)
})?;
Ok(Self { child })
}
}
fn missing_ffmpeg(error: std::io::Error, what: &str, alternative: &str) -> std::io::Error {
if error.kind() != std::io::ErrorKind::NotFound {
return error;
}
std::io::Error::new(
std::io::ErrorKind::NotFound,
format!(
"{what} (macOS: `brew install ffmpeg`, Debian/Ubuntu: `apt install ffmpeg`, \
conda: `conda install -c conda-forge ffmpeg`). {alternative}"
)
.trim_end()
.to_owned(),
)
}
fn wait_for_ffmpeg(child: &mut Child) -> std::io::Result<()> {
let status = child.wait()?;
if status.success() {
Ok(())
} else {
Err(std::io::Error::other(format!(
"ffmpeg exited with {status} — its error output is above"
)))
}
}
impl AnimationEncoder for FfmpegEncoder {
fn write_frame(&mut self, image: &Rgb8Image) -> std::io::Result<()> {
let stdin = self.child.stdin.as_mut().ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::BrokenPipe, "ffmpeg stdin was closed")
})?;
stdin.write_all(&image.pixels)
}
fn finish(mut self: Box<Self>) -> std::io::Result<()> {
drop(self.child.stdin.take());
wait_for_ffmpeg(&mut self.child)
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct VideoInfo {
pub width: usize,
pub height: usize,
pub fps: f64,
pub frames: Option<usize>,
}
impl VideoInfo {
pub fn probe(path: &Path) -> std::io::Result<Self> {
let output = Command::new("ffprobe")
.args(["-v", "error", "-select_streams", "v:0"])
.args(["-show_entries", "stream=width,height,r_frame_rate,nb_frames"])
.args(["-of", "csv=p=0"])
.arg(path)
.output()
.map_err(|error| missing_ffmpeg(error, "reading video needs ffmpeg on PATH", ""))?;
if !output.status.success() {
return Err(std::io::Error::other(format!(
"ffprobe could not read {}: {}",
path.display(),
String::from_utf8_lossy(&output.stderr).trim()
)));
}
Self::parse(&String::from_utf8_lossy(&output.stdout), path)
}
fn parse(text: &str, path: &Path) -> std::io::Result<Self> {
let malformed = || {
std::io::Error::other(format!(
"could not read the video stream of {} — is it a video file?",
path.display()
))
};
let line = text
.lines()
.find(|line| !line.trim().is_empty())
.ok_or_else(malformed)?;
let mut fields = line.trim().split(',');
let width: usize = fields
.next()
.ok_or_else(malformed)?
.trim()
.parse()
.map_err(|_| malformed())?;
let height: usize = fields
.next()
.ok_or_else(malformed)?
.trim()
.parse()
.map_err(|_| malformed())?;
let rate = fields.next().ok_or_else(malformed)?.trim();
let (num, den) = rate.split_once('/').unwrap_or((rate, "1"));
let num: f64 = num.parse().map_err(|_| malformed())?;
let den: f64 = den.parse().unwrap_or(1.0);
let fps = if den > 0.0 && num > 0.0 {
num / den
} else {
30.0
};
let frames = fields
.next()
.and_then(|field| field.trim().parse::<usize>().ok())
.filter(|&frames| frames > 0);
if width == 0 || height == 0 {
return Err(malformed());
}
Ok(Self {
width,
height,
fps,
frames,
})
}
}
pub struct FfmpegDecoder {
child: Child,
info: VideoInfo,
frame_bytes: usize,
buffer: Vec<u8>,
finished: bool,
}
impl FfmpegDecoder {
pub fn open(path: &Path, scale: Option<(usize, usize)>) -> std::io::Result<Self> {
let probed = VideoInfo::probe(path)?;
let info = match scale {
Some((width, height)) if width > 0 && height > 0 => VideoInfo {
width,
height,
fps: probed.fps,
frames: probed.frames,
},
_ => probed,
};
let mut command = Command::new("ffmpeg");
command
.args(["-hide_banner", "-loglevel", "error"])
.arg("-i")
.arg(path);
if scale.is_some() {
command.args(["-vf", &format!("scale={}:{}", info.width, info.height)]);
}
let child = command
.args(["-f", "rawvideo", "-pix_fmt", "rgb24", "-"])
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::inherit())
.spawn()
.map_err(|error| missing_ffmpeg(error, "reading video needs ffmpeg on PATH", ""))?;
Ok(Self {
child,
info,
frame_bytes: info.width * info.height * 3,
buffer: vec![0; info.width * info.height * 3],
finished: false,
})
}
pub fn info(&self) -> VideoInfo {
self.info
}
pub fn next_frame(&mut self) -> std::io::Result<Option<Rgb8Image>> {
if self.finished {
return Ok(None);
}
let stdout = self.child.stdout.as_mut().ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::BrokenPipe, "ffmpeg stdout was closed")
})?;
match read_exact_or_eof(stdout, &mut self.buffer[..self.frame_bytes])? {
true => Ok(Some(Rgb8Image {
width: self.info.width,
height: self.info.height,
pixels: self.buffer[..self.frame_bytes].to_vec(),
})),
false => {
self.finished = true;
wait_for_ffmpeg(&mut self.child)?;
Ok(None)
}
}
}
}
impl Drop for FfmpegDecoder {
fn drop(&mut self) {
if !self.finished {
let _ = self.child.kill();
let _ = self.child.wait();
}
}
}
fn read_exact_or_eof(reader: &mut impl std::io::Read, buffer: &mut [u8]) -> std::io::Result<bool> {
let mut filled = 0;
while filled < buffer.len() {
match reader.read(&mut buffer[filled..]) {
Ok(0) => return Ok(false),
Ok(n) => filled += n,
Err(error) if error.kind() == std::io::ErrorKind::Interrupted => {}
Err(error) => return Err(error),
}
}
Ok(true)
}
fn to_io_error<E: std::fmt::Display>(error: E) -> std::io::Error {
std::io::Error::other(error.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::viz::Rgb8Image;
fn frame(width: usize, height: usize, shade: u8) -> Rgb8Image {
Rgb8Image {
width,
height,
pixels: vec![shade; width * height * 3],
}
}
fn temp_path(name: &str) -> std::path::PathBuf {
let mut path = std::env::temp_dir();
path.push(format!(
"eventcv-video-test-{}-{}",
std::process::id(),
name
));
path
}
#[test]
fn format_is_read_from_the_extension() {
let cases = [
("a.gif", Some(AnimationFormat::Gif)),
("a.GIF", Some(AnimationFormat::Gif)),
("a.png", Some(AnimationFormat::Apng)),
("a.apng", Some(AnimationFormat::Apng)),
("a.mp4", Some(AnimationFormat::Mp4)),
("a.mov", Some(AnimationFormat::Mp4)),
("a.txt", None),
("a", None),
];
for (name, expected) in cases {
assert_eq!(
AnimationFormat::from_path(Path::new(name)),
expected,
"{name}"
);
}
}
#[test]
fn fps_clamps_and_converts() {
assert_eq!(Fps::new(f64::NAN).get(), 30.0);
assert_eq!(Fps::new(0.0).get(), 0.1);
assert_eq!(Fps::new(1e9).get(), 1000.0);
assert_eq!(Fps::new(100.0).centiseconds(), 1); assert_eq!(Fps::new(50.0).centiseconds(), 2);
assert_eq!(Fps::new(10.0).centiseconds(), 10);
}
#[test]
fn apng_writes_a_multi_frame_file() {
let path = temp_path("apng.png");
let mut encoder: Box<dyn AnimationEncoder> =
Box::new(ApngEncoder::new(&path, 3, 4, 4, Fps::new(10.0)).unwrap());
for shade in [0u8, 128, 255] {
encoder.write_frame(&frame(4, 4, shade)).unwrap();
}
encoder.finish().unwrap();
let bytes = std::fs::read(&path).unwrap();
assert_eq!(&bytes[..8], b"\x89PNG\r\n\x1a\n");
assert!(bytes.windows(4).any(|w| w == b"acTL"));
assert!(bytes.windows(4).any(|w| w == b"fcTL"));
std::fs::remove_file(&path).ok();
}
#[test]
fn gif_writes_a_multi_frame_file() {
let path = temp_path("gif.gif");
let mut encoder: Box<dyn AnimationEncoder> =
Box::new(GifEncoder::new(&path, 4, 4, Fps::new(10.0)).unwrap());
for shade in [0u8, 128, 255] {
encoder.write_frame(&frame(4, 4, shade)).unwrap();
}
encoder.finish().unwrap();
let bytes = std::fs::read(&path).unwrap();
assert_eq!(&bytes[..6], b"GIF89a");
assert_eq!(bytes.last(), Some(&0x3B)); std::fs::remove_file(&path).ok();
}
#[test]
fn video_info_parses_ffprobe_csv() {
let path = Path::new("clip.mp4");
let info = VideoInfo::parse("64,48,30/1\n", path).unwrap();
assert_eq!((info.width, info.height), (64, 48));
assert!((info.fps - 30.0).abs() < 1e-9);
let ntsc = VideoInfo::parse("1920,1080,30000/1001", path).unwrap();
assert!((ntsc.fps - 29.97).abs() < 0.01);
assert!((VideoInfo::parse("8,8,25", path).unwrap().fps - 25.0).abs() < 1e-9);
}
#[test]
fn video_info_rejects_nonsense() {
let path = Path::new("notes.txt");
for text in ["", "\n", "not,a,video", "0,0,30/1"] {
assert!(VideoInfo::parse(text, path).is_err(), "{text:?}");
}
}
#[test]
fn decoder_reads_back_every_frame_it_was_given() {
if Command::new("ffmpeg").arg("-version").output().is_err() {
return; }
let path = temp_path("roundtrip.mp4");
let mut encoder: Box<dyn AnimationEncoder + Send> =
Box::new(FfmpegEncoder::new(&path, 32, 24, Fps::new(10.0)).unwrap());
for shade in [0u8, 60, 120, 180, 240] {
encoder.write_frame(&frame(32, 24, shade)).unwrap();
}
encoder.finish().unwrap();
let mut decoder = FfmpegDecoder::open(&path, None).unwrap();
assert_eq!((decoder.info().width, decoder.info().height), (32, 24));
let mut decoded = 0;
while let Some(image) = decoder.next_frame().unwrap() {
assert_eq!(image.pixels.len(), 32 * 24 * 3);
decoded += 1;
}
assert_eq!(decoded, 5);
assert!(decoder.next_frame().unwrap().is_none());
std::fs::remove_file(&path).ok();
}
#[test]
fn decoder_can_scale_on_the_way_out() {
if Command::new("ffmpeg").arg("-version").output().is_err() {
return;
}
let path = temp_path("scaled.mp4");
let mut encoder: Box<dyn AnimationEncoder + Send> =
Box::new(FfmpegEncoder::new(&path, 64, 64, Fps::new(10.0)).unwrap());
encoder.write_frame(&frame(64, 64, 128)).unwrap();
encoder.finish().unwrap();
let mut decoder = FfmpegDecoder::open(&path, Some((16, 16))).unwrap();
let image = decoder.next_frame().unwrap().expect("one frame");
assert_eq!((image.width, image.height), (16, 16));
assert_eq!(image.pixels.len(), 16 * 16 * 3);
std::fs::remove_file(&path).ok();
}
#[test]
fn read_exact_or_eof_reports_a_clean_end() {
let mut full = [0u8; 4];
assert!(read_exact_or_eof(&mut &b"abcd"[..], &mut full).unwrap());
assert_eq!(&full, b"abcd");
assert!(!read_exact_or_eof(&mut &b""[..], &mut full).unwrap());
assert!(!read_exact_or_eof(&mut &b"ab"[..], &mut full).unwrap());
}
#[test]
fn unknown_extension_is_rejected_with_a_useful_message() {
let error = encoder_for(Path::new("out.avi"), 1, 4, 4, Fps::default())
.err()
.expect("an unknown extension must not open an encoder");
assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
assert!(error.to_string().contains(".gif"));
}
#[test]
fn missing_ffmpeg_names_the_fix() {
if Command::new("ffmpeg").arg("-version").output().is_err() {
let error = FfmpegEncoder::new(&temp_path("x.mp4"), 4, 4, Fps::default())
.err()
.expect("spawning ffmpeg must fail when it is not installed");
assert_eq!(error.kind(), std::io::ErrorKind::NotFound);
assert!(error.to_string().contains("ffmpeg"));
assert!(error.to_string().contains(".gif"));
}
}
}