Skip to main content

ez_ffmpeg/core/frame_export/
video.rs

1//! The [`FrameExtractor`] builder and its `frames()` / `collect_frames()` runs.
2
3use super::error::FrameExportError;
4use super::frame::VideoFrame;
5use super::guard::ColorGuard;
6use super::iter::FrameIter;
7use super::options::{ColorPolicy, ConversionPrecision, PixelLayout, Sampling};
8use super::resolve::{resolve_and_build_desc, ResolvePlan, UniformResolve};
9use super::sampler::{ExportSampler, UniformSpan};
10use super::selector::InputSelector;
11use super::sink::ExportSink;
12use crate::core::context::demuxer::Demuxer;
13use crate::core::context::ffmpeg_context::FfmpegContext;
14use crate::core::context::filter_complex::FilterComplex;
15use crate::core::context::input::Input;
16use crate::core::context::output::{Output, VSyncMethod};
17use crate::core::filter::frame_pipeline_builder::FramePipelineBuilder;
18use crate::core::scheduler::ffmpeg_scheduler::FfmpegScheduler;
19use ffmpeg_sys_next::AVMediaType::AVMEDIA_TYPE_VIDEO;
20
21/// Extracts packed RGB/gray frames from a video input, one pass, for AI/CV.
22///
23/// Build one with [`FrameExtractor::new`], chain options, then call
24/// [`frames`](FrameExtractor::frames) for a streaming iterator or
25/// [`collect_frames`](FrameExtractor::collect_frames) for a `Vec`.
26pub struct FrameExtractor {
27    input: Input,
28    sampling: Sampling,
29    video_stream_index: Option<usize>,
30    width: Option<u32>,
31    height: Option<u32>,
32    pixel: PixelLayout,
33    color: ColorPolicy,
34    precision: ConversionPrecision,
35    max_frames: Option<u64>,
36    start_time_us: Option<i64>,
37    duration_us: Option<i64>,
38    duration_hint_us: Option<i64>,
39    channel_capacity: usize,
40}
41
42impl FrameExtractor {
43    /// Creates an extractor over `input` (a path, URL, or anything convertible
44    /// into an [`Input`]). Defaults: all frames, source size, `Rgb24`,
45    /// [`ColorPolicy::Tagged`], channel capacity 1.
46    pub fn new(input: impl Into<Input>) -> Self {
47        Self {
48            input: input.into(),
49            sampling: Sampling::All,
50            video_stream_index: None,
51            width: None,
52            height: None,
53            pixel: PixelLayout::Rgb24,
54            color: ColorPolicy::Tagged,
55            precision: ConversionPrecision::Standard,
56            max_frames: None,
57            start_time_us: None,
58            duration_us: None,
59            duration_hint_us: None,
60            channel_capacity: 1,
61        }
62    }
63
64    /// Sets the sampling strategy (default [`Sampling::All`]).
65    pub fn sampling(mut self, sampling: Sampling) -> Self {
66        self.sampling = sampling;
67        self
68    }
69
70    /// Selects an explicit video stream by absolute index (default: best video
71    /// stream). Also pins the input-side per-frame HDR guard / color stamp to
72    /// that stream — without it they bind to the first video stream, which
73    /// coincides with the exported one except on exotic multi-video layouts.
74    pub fn video_stream_index(mut self, index: usize) -> Self {
75        self.video_stream_index = Some(index);
76        self
77    }
78
79    /// Sets the output width in pixels; height (if unset) is derived keeping the
80    /// aspect ratio (`scale=W:-2`).
81    pub fn width(mut self, width: u32) -> Self {
82        self.width = Some(width);
83        self
84    }
85
86    /// Sets the output height in pixels; width (if unset) is derived keeping the
87    /// aspect ratio (`scale=-2:H`).
88    pub fn height(mut self, height: u32) -> Self {
89        self.height = Some(height);
90        self
91    }
92
93    /// Sets the packed pixel layout (default [`PixelLayout::Rgb24`]).
94    pub fn pixel(mut self, layout: PixelLayout) -> Self {
95        self.pixel = layout;
96        self
97    }
98
99    /// Sets the color-interpretation policy (default [`ColorPolicy::Tagged`]).
100    pub fn color(mut self, policy: ColorPolicy) -> Self {
101        self.color = policy;
102        self
103    }
104
105    /// Sets the swscale precision tier for the conversion stage — the single
106    /// swscale pass that runs any resize plus the pixel-format conversion
107    /// (default [`ConversionPrecision::Standard`], which matches the FFmpeg
108    /// CLI's default scaler flags). [`ConversionPrecision::High`] opts into
109    /// accurate rounding + full chroma interpolation at a several-fold
110    /// conversion cost — see the enum docs for the tradeoff.
111    pub fn conversion_precision(mut self, precision: ConversionPrecision) -> Self {
112        self.precision = precision;
113        self
114    }
115
116    /// Caps the number of exported frames. The sink owns this exact count.
117    pub fn max_frames(mut self, max: u64) -> Self {
118        self.max_frames = Some(max);
119        self
120    }
121
122    /// Seeks to this start time (microseconds) before extracting.
123    pub fn start_time_us(mut self, us: i64) -> Self {
124        self.start_time_us = Some(us);
125        self
126    }
127
128    /// Limits extraction to this many microseconds of content past the start.
129    pub fn duration_us(mut self, us: i64) -> Self {
130        self.duration_us = Some(us);
131        self
132    }
133
134    /// Supplies a duration (microseconds) for `UniformN` when the input is not
135    /// probeable (live/piped). Ignored by other sampling modes. Takes precedence
136    /// over the container/stream duration, but not over an explicit
137    /// [`duration_us`](FrameExtractor::duration_us) trim window.
138    pub fn duration_hint_us(mut self, us: i64) -> Self {
139        self.duration_hint_us = Some(us);
140        self
141    }
142
143    /// Sets the prefetch channel capacity (default 1, minimum 1). Each slot
144    /// costs one packed frame of memory.
145    pub fn channel_capacity(mut self, capacity: usize) -> Self {
146        self.channel_capacity = capacity;
147        self
148    }
149
150    /// Starts the run and returns a streaming iterator over the exported frames.
151    ///
152    /// Option and stream/HDR resolution errors surface here, before any frame is
153    /// produced; runtime failures surface as the iterator's terminal `Err`.
154    pub fn frames(self) -> crate::error::Result<FrameIter> {
155        self.validate()?;
156
157        let capacity = self.channel_capacity.max(1);
158        let (tx, rx) = crossbeam_channel::bounded(capacity);
159
160        // Output-side sink on the single graph output stream ([export] => 0).
161        let sink = ExportSink::new(tx, self.sampling, self.pixel, self.max_frames);
162        let sink_pipeline = FramePipelineBuilder::new(AVMEDIA_TYPE_VIDEO)
163            .filter("frame_export_sink", Box::new(sink))
164            .set_stream_index(0)
165            .build();
166
167        // UniformN selects frames pre-filtergraph on the input side; the grid
168        // span is published through this cell at open time (§4.4) and read by
169        // the sampler at its first frame.
170        let uniform_n = match self.sampling {
171            Sampling::UniformN(n) => Some(n),
172            _ => None,
173        };
174        let span_cell: UniformSpan = std::sync::Arc::new(std::sync::OnceLock::new());
175
176        // Input decoder fast path + time window. The extractor's own options
177        // take precedence, but a start/duration preconfigured directly on the
178        // Input must feed the SAME machinery (trim boundary, span sizing) —
179        // otherwise it would reintroduce the GOP lead-in mis-anchoring the
180        // boundary exists to prevent.
181        let mut input = self.input;
182        let effective_start = self.start_time_us.or(input.start_time_us);
183        let effective_duration = self.duration_us.or(input.recording_time_us);
184        if matches!(self.sampling, Sampling::KeyframesOnly) {
185            input = input.set_video_codec_opt("skip_frame", "nokey");
186        }
187        if let Some(start) = self.start_time_us {
188            input = input.set_start_time_us(start);
189        }
190        if let Some(dur) = self.duration_us {
191            input = input.set_recording_time_us(dur);
192        } else if uniform_n.is_some() && input.recording_time_us.is_none() {
193            if let Some(hint) = self.duration_hint_us {
194                // The hint declares the sampled span; also use it as the demux
195                // stop boundary so an unbounded (live/piped) input terminates
196                // once the grid is covered instead of decoding forever. A
197                // recording time already set on the Input wins over the hint.
198                input = input.set_recording_time_us(hint);
199            }
200        }
201        // Input-side pipeline, every run: the color guard re-checks each
202        // decoded (pre-conversion) frame for HDR and applies the
203        // TaggedOrResolutionGuess stamp; UniformN chains its sampler after the
204        // guard so selection only ever sees guarded, stamped frames.
205        let mut builder = FramePipelineBuilder::new(AVMEDIA_TYPE_VIDEO).filter(
206            "frame_export_color_guard",
207            Box::new(ColorGuard::new(self.color)),
208        );
209        // With a start time (from the extractor or preconfigured on the
210        // Input) the demux timeline is re-zeroed at the request and the
211        // in-graph trim drops pts < 0 — but the container seek lands on a
212        // keyframe at or BEFORE the request, so on GOP video this pipeline
213        // sees negative-pts lead-in first. The boundary makes the input-side
214        // sampler/selector skip that lead-in instead of anchoring grids on
215        // (or spending selections on) frames the trim will destroy.
216        let trim_boundary = effective_start.map(|_| 0i64);
217        if let Some(n) = uniform_n {
218            let sampler = ExportSampler::new(n, span_cell.clone(), trim_boundary);
219            builder = builder.filter("frame_export_sampler", Box::new(sampler));
220        } else if let Some(selector) = InputSelector::for_sampling(&self.sampling, trim_boundary) {
221            // EveryNth/EverySec drop unselected frames before the filtergraph,
222            // so they never pay the scale/format conversion.
223            builder = builder.filter("frame_export_selector", Box::new(selector));
224        }
225        // Bind to the explicit stream when given; otherwise the video stream
226        // by media type (the resolver's best-stream selection coincides for
227        // single-video-stream inputs, the common case — the resolver rejects
228        // or warns on the mismatching layouts).
229        if let Some(idx) = self.video_stream_index {
230            builder = builder.set_stream_index(idx);
231        }
232        input = input.add_frame_pipeline(builder.build());
233
234        // Null output: vsync passthrough (S4) preserves 1:1 frames and PTS.
235        let mut output = Output::from("-")
236            .set_format("null")
237            .set_vsync_method(VSyncMethod::VsyncPassthrough)
238            .add_stream_map("[export]")
239            .add_frame_pipeline(sink_pipeline);
240        // UniformN emits <= n unique frames and must reach EOF for its flush, so
241        // it gets no encoder terminator; the sink's u64 cap stays authoritative.
242        if uniform_n.is_none() {
243            if let Some(max) = self.max_frames {
244                // Upstream terminator only; the sink owns the exact public cap
245                // (S5). If the cap exceeds i64, skip the terminator (the sink's
246                // u64 cap stays authoritative) rather than wrapping negative.
247                if let Ok(max_i64) = i64::try_from(max) {
248                    output = output.set_max_video_frames(Some(max_i64));
249                }
250            }
251        }
252
253        // Deferred resolver: runs against the opened demuxer (S8/S9).
254        let resolve_plan = ResolvePlan {
255            stream_index: self.video_stream_index,
256            width: self.width,
257            height: self.height,
258            pixel: self.pixel,
259            color: self.color,
260            precision: self.precision,
261            input_side_sampling: matches!(
262                self.sampling,
263                Sampling::UniformN(_) | Sampling::EveryNth(_) | Sampling::EverySec(_)
264            ),
265            uniform: uniform_n.map(|_| UniformResolve {
266                span_cell: span_cell.clone(),
267                duration_hint_us: self.duration_hint_us,
268                duration_us: effective_duration,
269                start_time_us: effective_start,
270            }),
271        };
272        let resolver = move |demuxs: &[Demuxer]| -> crate::error::Result<FilterComplex> {
273            let demux = demuxs.first().ok_or(FrameExportError::NoVideoStream)?;
274            // SAFETY: the demuxer is opened for the lifetime of this call.
275            let desc = unsafe { resolve_and_build_desc(demux.in_fmt_ctx_ptr(), &resolve_plan)? };
276            Ok(desc.into())
277        };
278
279        let context = FfmpegContext::builder()
280            .input(input)
281            .output(output)
282            .add_deferred_filter_desc(Box::new(resolver))
283            .build()?;
284        let scheduler = FfmpegScheduler::new(context).start()?;
285        Ok(FrameIter::new(rx, scheduler))
286    }
287
288    /// Runs to completion and collects every exported frame into a `Vec`.
289    ///
290    /// On a terminal error the partial frames are dropped and the error is
291    /// returned; use [`frames`](FrameExtractor::frames) to keep partial output.
292    pub fn collect_frames(self) -> crate::error::Result<Vec<VideoFrame>> {
293        self.frames()?.collect()
294    }
295
296    fn validate(&self) -> crate::error::Result<()> {
297        if self.channel_capacity == 0 {
298            return Err(invalid("channel_capacity must be >= 1"));
299        }
300        match self.sampling {
301            Sampling::EveryNth(0) => {
302                return Err(invalid("EveryNth requires n >= 1"));
303            }
304            Sampling::EverySec(s) if !s.is_finite() || s <= 0.0 => {
305                return Err(invalid(
306                    "EverySec requires a finite, positive seconds value",
307                ));
308            }
309            Sampling::UniformN(0) => {
310                return Err(invalid("UniformN requires n >= 1"));
311            }
312            _ => {}
313        }
314        if self.width == Some(0) {
315            return Err(invalid("width must be > 0"));
316        }
317        if self.height == Some(0) {
318            return Err(invalid("height must be > 0"));
319        }
320        if self.max_frames == Some(0) {
321            return Err(invalid("max_frames must be > 0"));
322        }
323        // Effective value: a recording time preconfigured on the Input feeds
324        // the same trim machinery, where 0 means "no limit" to FFmpeg — the
325        // opposite of the documented meaning.
326        if let Some(d) = self.duration_us.or(self.input.recording_time_us) {
327            if d <= 0 {
328                return Err(invalid(
329                    "duration_us (or the Input's recording time) must be > 0",
330                ));
331            }
332        }
333        // The hint is only consumed by UniformN; other modes ignore it as
334        // documented, so it is validated only where it has meaning.
335        if matches!(self.sampling, Sampling::UniformN(_)) {
336            if let Some(h) = self.duration_hint_us {
337                if h <= 0 {
338                    return Err(invalid("duration_hint_us must be > 0"));
339                }
340            }
341        }
342        Ok(())
343    }
344}
345
346fn invalid(msg: &str) -> crate::error::Error {
347    FrameExportError::InvalidOption(msg.to_string()).into()
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353
354    #[test]
355    fn zero_every_nth_is_rejected() {
356        let r = FrameExtractor::new("x.mp4")
357            .sampling(Sampling::EveryNth(0))
358            .validate();
359        assert!(matches!(
360            r,
361            Err(crate::error::Error::FrameExport(
362                FrameExportError::InvalidOption(_)
363            ))
364        ));
365    }
366
367    #[test]
368    fn non_finite_every_sec_is_rejected() {
369        assert!(FrameExtractor::new("x.mp4")
370            .sampling(Sampling::EverySec(f64::NAN))
371            .validate()
372            .is_err());
373        assert!(FrameExtractor::new("x.mp4")
374            .sampling(Sampling::EverySec(-1.0))
375            .validate()
376            .is_err());
377        assert!(FrameExtractor::new("x.mp4")
378            .sampling(Sampling::EverySec(0.0))
379            .validate()
380            .is_err());
381    }
382
383    #[test]
384    fn zero_dimensions_and_capacity_rejected() {
385        assert!(FrameExtractor::new("x.mp4").width(0).validate().is_err());
386        assert!(FrameExtractor::new("x.mp4").height(0).validate().is_err());
387        assert!(FrameExtractor::new("x.mp4")
388            .channel_capacity(0)
389            .validate()
390            .is_err());
391        assert!(FrameExtractor::new("x.mp4")
392            .max_frames(0)
393            .validate()
394            .is_err());
395    }
396
397    #[test]
398    fn defaults_validate() {
399        assert!(FrameExtractor::new("x.mp4").validate().is_ok());
400        assert!(FrameExtractor::new("x.mp4")
401            .sampling(Sampling::EverySec(1.5))
402            .width(224)
403            .height(224)
404            .max_frames(16)
405            .validate()
406            .is_ok());
407    }
408}