1use std::path::Path;
2use std::sync::atomic::{AtomicBool, Ordering};
3use std::sync::mpsc;
4
5use super::export::{ExportOptions, Exported, MAX_FPS, Progress, Stage, WRITTEN};
6use super::mux::{
7 RATE, Shape, Sink, Spill, VEnc, aac_encoder, check_at, encoder_for, joining, seal,
8};
9use super::picture::Blender;
10use super::route::{Caps, Choice, Picked, Prefs, pick_encode};
11use super::video::VideoFrame;
12use crate::media::{AudioBuffer, Codec, Error, Frame, PixelFormat, Result, Route, TimeBase};
13
14const AUDIO_STOPPED: &str = "the audio stopped";
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub struct Canvas {
18 pub width: u32,
19 pub height: u32,
20 pub fps: (u32, u32),
21 pub frames: Option<u64>,
22}
23
24fn samples_at(frames: u64, fps: (u32, u32)) -> u64 {
25 let num = frames as u128 * RATE as u128 * fps.1 as u128;
26 let den = (fps.0 as u128).max(1);
27 u64::try_from((2 * num + den) / (2 * den)).unwrap_or(u64::MAX)
28}
29
30fn native_prefs(prefs: &Prefs, reasons: &mut Vec<String>) -> Prefs {
31 let encode = match &prefs.encode {
32 Choice::Ffmpeg(_) => {
33 reasons.push(
34 "frames made by the caller are encoded by our own encoders, not ffmpeg".to_string(),
35 );
36 Choice::Auto
37 }
38 c => c.clone(),
39 };
40 Prefs {
41 decode: Choice::Auto,
42 encode,
43 hevc_encode_fallback: prefs.hevc_encode_fallback.clone(),
44 }
45}
46
47fn open_encoder(shape: Shape, video: &mut Picked, crf: u8) -> Result<VEnc> {
48 match encoder_for(shape, video, crf, false, 0, 0) {
49 Ok(e) => Ok(e),
50 Err(e) if video.route == Route::Gpu => {
51 video.reasons.push(format!(
52 "the GPU encoder failed ({e}); encoding H.264 on the CPU instead"
53 ));
54 video.route = Route::Cpu;
55 video.codec = Codec::H264;
56 encoder_for(shape, video, crf, false, 0, 0)
57 }
58 Err(e) => Err(e),
59 }
60}
61
62fn checked(n: u64, frame: &Frame, size: (u32, u32)) -> Result<()> {
63 if (frame.width, frame.height) != size {
64 return Err(Error::Invalid(format!(
65 "frame {n} is {}x{}, the export is {}x{}",
66 frame.width, frame.height, size.0, size.1
67 )));
68 }
69 if frame.format != PixelFormat::Yuv420p8 {
70 return Err(Error::Invalid(format!(
71 "frame {n} is {:?}, the export takes 8-bit 4:2:0",
72 frame.format
73 )));
74 }
75 crate::filter::check_frame(frame).map_err(|e| Error::Invalid(format!("frame {n}: {e}")))
76}
77
78pub fn export_frames(
79 out: &Path,
80 canvas: Canvas,
81 mut opts: ExportOptions,
82 mut frames: impl FnMut(u64, Option<Frame>) -> Result<Option<VideoFrame>>,
83 mut audio: impl FnMut(u64, usize) -> Result<AudioBuffer> + Send,
84 mut progress: impl FnMut(&Progress) -> bool,
85) -> Result<Exported> {
86 let size = (canvas.width, canvas.height);
87 let fps = canvas.fps;
88 if size.0 == 0 || size.1 == 0 || !size.0.is_multiple_of(2) || !size.1.is_multiple_of(2) {
89 return Err(Error::Invalid(format!(
90 "a picture of {}x{}; the width and height must be even",
91 size.0, size.1
92 )));
93 }
94 if fps.0 == 0 || fps.1 == 0 || fps.0 / fps.1 > MAX_FPS {
95 return Err(Error::Invalid(format!(
96 "a frame rate of {}/{}",
97 fps.0, fps.1
98 )));
99 }
100 if !matches!(opts.video, Codec::H264 | Codec::Hevc) {
101 return Err(Error::Unsupported(format!(
102 "frames export to H.264 or HEVC, not {}",
103 opts.video
104 )));
105 }
106 let mut reasons = Vec::new();
107 let prefs = native_prefs(&opts.prefs, &mut reasons);
108 let caps = Caps::detect(&prefs);
109 let mut video = pick_encode(opts.video, &prefs, &caps)?;
110 let shape = (size, fps, PixelFormat::Yuv420p8);
111 let mut enc = open_encoder(shape, &mut video, opts.crf)?;
112 for r in &video.reasons {
113 if !reasons.contains(r) {
114 reasons.push(r.clone());
115 }
116 }
117 let work = std::path::absolute(out.with_extension("parts"))?;
118 if work.exists() {
119 std::fs::remove_dir_all(&work).map_err(|e| Error::at(&work, e))?;
120 }
121 std::fs::create_dir_all(&work).map_err(|e| Error::at(&work, e))?;
122 let part = work.join("joined.mp4");
123 let route = video.route;
124 let result = (|| -> Result<u64> {
125 let mut aac = aac_encoder(opts.audio_kbps)?;
126 let timescale = fps.0.saturating_mul(64);
127 let fd = fps.1 as i64 * 64;
128 let mut sink = Sink::new(video.codec, size, 0, (timescale, fd), &part);
129 let stop = AtomicBool::new(false);
130 let batch = u64::from((fps.0 / fps.1).max(1));
131 let (tx, rx) = mpsc::sync_channel::<u64>(4);
132 let (count, sound) = std::thread::scope(|s| -> Result<_> {
133 let (aac, stop, work) = (&mut aac, &stop, &work);
134 let sound = s.spawn(move || -> Result<(Spill, u64)> {
135 let mut packets = Spill::create(work, "premixed.aac")?;
136 let mut first = 0u64;
137 for upto in rx {
138 if stop.load(Ordering::Relaxed) {
139 return Err(Error::Cancelled);
140 }
141 let end = samples_at(upto, fps);
142 let n = usize::try_from(end.saturating_sub(first))
143 .map_err(|_| Error::Invalid("too much audio for one call".into()))?;
144 let buf = audio(upto, n)?;
145 check_at(&format_args!("the audio to frame {upto}"), "it", n, &buf)?;
146 packets.extend(aac.encode(&buf)?)?;
147 first = end;
148 }
149 if stop.load(Ordering::Relaxed) {
150 return Err(Error::Cancelled);
151 }
152 packets.extend(aac.flush())?;
153 Ok((packets, first))
154 });
155 let mut pictures = || -> Result<u64> {
156 let time_base = TimeBase::new(fps.1, fps.0);
157 let mut blender = Blender::new();
158 let mut spent: Option<Frame> = None;
159 let mut n = 0u64;
160 let tell = |n: u64, progress: &mut dyn FnMut(&Progress) -> bool| {
161 let done = match canvas.frames {
162 Some(t) if t > 0 => (n as f32 / t as f32).min(1.0) * WRITTEN,
163 _ => 0.0,
164 };
165 let keep = progress(&Progress {
166 stage: Stage::Encoding,
167 done,
168 encoder: route,
169 });
170 if keep { Ok(()) } else { Err(Error::Cancelled) }
171 };
172 tell(0, &mut progress)?;
173 while canvas.frames.is_none_or(|t| n < t)
174 && let Some(frame) = frames(n, spent.take())?
175 {
176 let seconds = n as f64 * fps.1 as f64 / fps.0 as f64;
177 let layer = match opts.overlay.as_mut() {
178 Some(o) => o(seconds, size.0, size.1),
179 None => None,
180 };
181 let pts = i64::try_from(n).unwrap_or(i64::MAX);
182 let packets = match frame {
183 VideoFrame::Gpu(mut g) if layer.is_none() && enc.on_gpu() => {
184 if g.size() != size {
185 return Err(Error::Invalid(format!(
186 "frame {n} is {}x{}, the export is {}x{}",
187 g.size().0,
188 g.size().1,
189 size.0,
190 size.1
191 )));
192 }
193 g.retime(pts, time_base);
194 enc.encode_gpu(&g)?
195 }
196 other => {
197 let mut f = match other {
198 VideoFrame::Cpu(f) => f,
199 VideoFrame::Gpu(g) => g.download()?,
200 };
201 checked(n, &f, size)?;
202 f.pts = pts;
203 f.time_base = time_base;
204 if let Some(layer) = &layer {
205 blender.draw(&mut f, layer);
206 }
207 let p = enc.encode(&f)?;
208 spent = Some(f);
209 p
210 }
211 };
212 sink.feed(&enc, packets)?;
213 n += 1;
214 if n.is_multiple_of(batch) && tx.send(n).is_err() {
215 return Err(Error::Backend(AUDIO_STOPPED.to_string()));
216 }
217 if n.is_multiple_of(15) {
218 tell(n, &mut progress)?;
219 }
220 }
221 if n == 0 {
222 return Err(Error::Invalid("no frames to export".to_string()));
223 }
224 if !n.is_multiple_of(batch) && tx.send(n).is_err() {
225 return Err(Error::Backend(AUDIO_STOPPED.to_string()));
226 }
227 let packets = enc.flush()?;
228 sink.feed(&enc, packets)?;
229 Ok(n)
230 };
231 let made = pictures();
232 if made.is_err() {
233 stop.store(true, Ordering::Relaxed);
234 }
235 drop(tx);
236 let sound = sound
237 .join()
238 .unwrap_or_else(|_| Err(Error::Backend("the audio failed".to_string())));
239 match (made, sound) {
240 (Ok(n), Ok(sound)) => Ok((n, sound)),
241 (Err(Error::Backend(m)), Err(e)) if m == AUDIO_STOPPED => Err(e),
242 (Err(e), _) | (Ok(_), Err(e)) => Err(e),
243 }
244 })?;
245 let mux = sink.done()?;
246 if !progress(&joining(0.0, route)) {
247 return Err(Error::Cancelled);
248 }
249 let (packets, total) = sound;
250 seal(
251 mux,
252 &aac,
253 opts.audio_kbps,
254 (packets, Some(total)),
255 (&part, out),
256 &mut |f| progress(&joining(f, route)),
257 )?;
258 Ok(count)
259 })();
260 let _ = std::fs::remove_dir_all(&work);
261 let count = result?;
262 if !progress(&Progress {
263 stage: Stage::Joining,
264 done: 1.0,
265 encoder: route,
266 }) {
267 return Err(Error::Cancelled);
268 }
269 Ok(Exported {
270 path: out.to_path_buf(),
271 frames: count,
272 copied: 0,
273 encoded: count,
274 smart: false,
275 video,
276 reasons,
277 })
278}