ranim/cmd/render/
file_writer.rs

1use std::{
2    io::Write,
3    path::PathBuf,
4    process::{Child, ChildStdin, Command, Stdio},
5};
6
7use crate::OutputFormat;
8use tracing::info;
9
10/// Extension trait providing ffmpeg encoding parameters for [`OutputFormat`].
11pub(crate) trait OutputFormatExt {
12    /// Returns `(video_codec, pixel_format, file_extension)`.
13    fn encoding_params(&self) -> (&'static str, &'static str, &'static str);
14    /// Returns extra codec arguments for ffmpeg.
15    fn extra_args(&self) -> &'static [&'static str];
16    /// Whether this format has an alpha channel.
17    fn has_alpha(&self) -> bool;
18    /// Whether the `eq` video filter is compatible with this format.
19    fn supports_eq_filter(&self) -> bool;
20}
21
22impl OutputFormatExt for OutputFormat {
23    fn encoding_params(&self) -> (&'static str, &'static str, &'static str) {
24        match self {
25            Self::Mp4 => ("libx264", "yuv420p", "mp4"),
26            Self::Webm => ("libvpx-vp9", "yuva420p", "webm"),
27            Self::Mov => ("prores_ks", "yuva444p10le", "mov"),
28            Self::Gif => ("gif", "rgb8", "gif"),
29        }
30    }
31
32    fn extra_args(&self) -> &'static [&'static str] {
33        match self {
34            Self::Mov => &["-profile:v", "4444"],
35            _ => &[],
36        }
37    }
38
39    fn has_alpha(&self) -> bool {
40        matches!(self, Self::Webm | Self::Mov)
41    }
42
43    fn supports_eq_filter(&self) -> bool {
44        !self.has_alpha()
45    }
46}
47
48#[derive(Debug, Clone)]
49pub struct FileWriterBuilder {
50    pub file_path: PathBuf,
51    pub width: u32,
52    pub height: u32,
53    pub fps: u32,
54    pub vf_args: Vec<String>,
55
56    pub video_codec: String,
57    pub pixel_format: String,
58    pub extra_codec_args: Vec<String>,
59}
60
61impl Default for FileWriterBuilder {
62    fn default() -> Self {
63        Self {
64            file_path: PathBuf::from("output.mp4"),
65            width: 1920,
66            height: 1080,
67            fps: 60,
68
69            vf_args: vec!["eq=saturation=1.0:gamma=1.0".to_string()],
70            video_codec: "libx264".to_string(),
71            pixel_format: "yuv420p".to_string(),
72            extra_codec_args: Vec::new(),
73        }
74    }
75}
76
77#[allow(unused)]
78impl FileWriterBuilder {
79    pub fn with_file_path(mut self, file_path: PathBuf) -> Self {
80        self.file_path = file_path;
81        self
82    }
83
84    pub fn with_size(mut self, width: u32, height: u32) -> Self {
85        self.width = width;
86        self.height = height;
87        self
88    }
89
90    pub fn with_fps(mut self, fps: u32) -> Self {
91        self.fps = fps;
92        self
93    }
94
95    pub fn with_output_format(mut self, format: OutputFormat) -> Self {
96        let (codec, pix_fmt, ext) = format.encoding_params();
97        self.video_codec = codec.to_string();
98        self.pixel_format = pix_fmt.to_string();
99        self.extra_codec_args = format.extra_args().iter().map(|s| s.to_string()).collect();
100        // Update file extension to match the format
101        self.file_path = self.file_path.with_extension(ext);
102        // The eq filter doesn't support alpha pixel formats
103        if !format.supports_eq_filter() {
104            self.vf_args.clear();
105        }
106        // GIF timing uses centiseconds (10ms units), so fps above 50
107        // gets rounded and causes incorrect playback speed.
108        if format == OutputFormat::Gif && self.fps > 50 {
109            self.fps = 50;
110        }
111        self
112    }
113
114    pub fn enable_fast_encoding(mut self) -> Self {
115        self.video_codec = "libx264rgb".to_string();
116        self.pixel_format = "rgb32".to_string();
117        self
118    }
119
120    pub fn output_gif(mut self) -> Self {
121        // TODO: use palette to improve gif quality
122        self.file_path = self.file_path.with_file_name(format!(
123            "{}.gif",
124            self.file_path.file_stem().unwrap().to_string_lossy()
125        ));
126        self.fps = 30;
127        self.video_codec = "gif".to_string();
128        self.pixel_format = "rgb8".to_string();
129        self
130    }
131
132    pub fn build(self) -> FileWriter {
133        let parent = self.file_path.parent().unwrap();
134        if !parent.exists() {
135            std::fs::create_dir_all(parent).unwrap();
136        }
137
138        let mut command = if which::which("ffmpeg").is_ok() {
139            info!("using ffmpeg found from path env");
140            Command::new("ffmpeg")
141        } else {
142            info!("using ffmpeg from current working dir");
143            Command::new("./ffmpeg")
144        };
145
146        let size = format!("{}x{}", self.width, self.height);
147        let fps = self.fps.to_string();
148        let file_path = self.file_path.to_string_lossy().to_string();
149
150        // Input options (before -i)
151        command.args([
152            "-y", "-f", "rawvideo", "-s", &size, "-pix_fmt", "rgba", "-r", &fps, "-i", "-",
153        ]);
154        // Output options (before output file)
155        command.args(["-an", "-loglevel", "error", "-vcodec", &self.video_codec]);
156        command.args(&self.extra_codec_args);
157        command.args(["-pix_fmt", &self.pixel_format]);
158        if !self.vf_args.is_empty() {
159            let vf = self.vf_args.join(",");
160            command.args(["-vf", &vf]);
161        }
162        // Output file must be last
163        command.arg(&file_path);
164        command.stdin(Stdio::piped());
165
166        let mut child = command.spawn().expect("Failed to spawn ffmpeg");
167        FileWriter {
168            child_in: child.stdin.take(),
169            child,
170        }
171    }
172}
173
174pub struct FileWriter {
175    child: Child,
176    child_in: Option<ChildStdin>,
177}
178
179impl Drop for FileWriter {
180    fn drop(&mut self) {
181        self.child_in
182            .as_mut()
183            .unwrap()
184            .flush()
185            .expect("Failed to flush ffmpeg");
186        drop(self.child_in.take());
187        self.child.wait().expect("Failed to wait ffmpeg");
188    }
189}
190
191impl FileWriter {
192    // pub fn builder() -> FileWriterBuilder {
193    //     FileWriterBuilder::default()
194    // }
195
196    pub fn write_frame(&mut self, frame: &[u8]) {
197        self.child_in
198            .as_mut()
199            .unwrap()
200            .write_all(frame)
201            .expect("Failed to write frame");
202    }
203}
ranim/cmd/render/file_writer.rs

ranim/cmd/render/
file_writer.rs
