Skip to main content

frame_source/
lib.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use std::path::PathBuf;
5
6use strand_dynamic_frame::{DynamicFrame, DynamicFrameOwned};
7
8pub mod pv_tiff_stack;
9use pv_tiff_stack::TiffImage;
10use serde::{Deserialize, Serialize};
11pub mod fmf_source;
12mod h264_annexb_splitter;
13pub mod h264_poc;
14pub mod h264_source;
15pub mod mp4_source;
16mod opt_openh264_decoder;
17pub mod srt_reader;
18pub mod strand_cam_mkv_source;
19
20mod ntp_timestamp;
21
22#[derive(thiserror::Error, Debug)]
23pub enum Error {
24    #[error("IO error: {0}")]
25    Io(#[from] std::io::Error),
26    #[error("SRT file {path} is malformed starting around line {line}")]
27    SrtParseError { path: PathBuf, line: usize },
28    #[error("expected SPS not found")]
29    ExpectedSpsNotFound,
30    #[error("expected PPS not found")]
31    ExpectedPpsNotFound,
32    #[error("fmf file with not enough data")]
33    FmfWithNotEnoughData,
34    #[error("JSON parse error")]
35    JsonParseError,
36    #[error("expected tiff image")]
37    ExpectedTiffImage,
38    #[error("no files found with pattern")]
39    NoFilesFound,
40    #[error("unsupported for estimating luminance range")]
41    UnsupportedForEsimatingLuminangeRange,
42    #[error("imagej data expected to be bytes")]
43    ImageJDataExpectedToBeBytes,
44    #[error("failed to read metadata")]
45    FailedToReadMetadata,
46    #[error("exif metadata does not start with expected magic string")]
47    ExifMetadataFailsMagic,
48    #[error("Skipping frames with H264 file is not supported.")]
49    SkippingFramesNotSupported,
50    #[error("Not implemented: {0}")]
51    NotImplemented(&'static str),
52    #[error("Requested SRT file as timestamp source, but no .srt file path given.")]
53    NoSrtPathGiven,
54    #[error("H264Error: {0}")]
55    H264Error(&'static str),
56    #[error("unexpected error reading NAL unit {nal_location_index} SEI: {e:?}")]
57    H264Nal {
58        nal_location_index: usize,
59        e: h264_reader::rbsp::BitReaderError,
60    },
61    #[error("PPS error {0}")]
62    H264Pps(String),
63    #[error("H264 timestamp error {0}")]
64    H264TimestampError(String),
65    #[error("H264 POC error {0}")]
66    H264Poc(String),
67    #[error("H264 UDU error {0}")]
68    UduError(String),
69    #[error("unexpected payload length")]
70    UnexpectedPayloadLength,
71    #[error("unexpected start code emulation prevention byte")]
72    UnexpectedStartCodeByte,
73    #[error("MP4 source error: {0}")]
74    Mp4SourceError(#[from] mp4_source::Mp4SourceError),
75    #[error("strand camera MKV source error: {0}")]
76    StrandMkvSourceError(#[from] strand_cam_mkv_source::StrandMkvSourceError),
77    #[error("srt file given, but not supported for this file type")]
78    NoSrtSupportForFileType,
79    #[error("unsupported option")]
80    UnsupportedOption,
81    #[error("input {0} is a file, but the extension was not recognized.")]
82    UnknownExtensionForFile(PathBuf),
83    #[error(
84        "Attempting to open \"{0}\" as directory with TIFF stack failed because it is not a directory."
85    )]
86    TiffStackNotDir(PathBuf),
87    #[error("{0}")]
88    FmfError(#[from] fmf::FMFError),
89    #[error("{0}")]
90    PatternError(#[from] glob::PatternError),
91    #[error("{0}")]
92    GlobError(#[from] glob::GlobError),
93    #[error("{0}")]
94    OutOfRangeError(#[from] chrono::OutOfRangeError),
95    #[error("{0}")]
96    ChronoParseError(#[from] chrono::ParseError),
97    #[error("{0}")]
98    TiffError(#[from] tiff::TiffError),
99    #[error("{0}")]
100    TryFromIntError(#[from] std::num::TryFromIntError),
101    #[error("{0}")]
102    ParseIntError(#[from] std::num::ParseIntError),
103    #[error("{0}")]
104    ExifError(#[from] exif::Error),
105    #[error("{0}")]
106    FromUtf8Error(#[from] std::string::FromUtf8Error),
107    #[error("{0}")]
108    SerdeJsonError(#[from] serde_json::Error),
109    #[error("{0}")]
110    MkvStrandError(#[from] mkv_strand_reader::Error),
111    #[cfg(feature = "openh264")]
112    #[error("OpenH264Error: {0}")]
113    OpenH264Error(#[from] openh264::Error),
114    #[cfg(feature = "openh264")]
115    #[error("OpenH264 failed to decode H.264 frame {frame}: {source}{hint}")]
116    H264DecodeFailed {
117        frame: usize,
118        /// Extra diagnostic context appended to the message, e.g. a note that
119        /// the stream's chroma subsampling is unsupported. Empty when there is
120        /// no likely-cause hint to offer.
121        hint: String,
122        #[source]
123        source: openh264::Error,
124    },
125    #[error("Mp4Error: {0}")]
126    Mp4Error(#[from] mp4::Error),
127    #[error("PreParserError: {0}")]
128    PreParserError(eyre::Report),
129}
130
131pub type Result<T> = std::result::Result<T, Error>;
132
133#[cfg(feature = "openh264")]
134pub const COMPILED_WITH_OPENH264: bool = true;
135#[cfg(not(feature = "openh264"))]
136pub const COMPILED_WITH_OPENH264: bool = false;
137
138/// A source of FrameData
139///
140/// The `frame0_time` method return value is an `Option` because we want to be
141/// able to parse sources without an absolute time for the first frame, such as
142/// normal MP4 video files. Similarly, we do not have a `len` method indicating
143/// number of frames because some sources (e.g. an .h264 file) do not store how
144/// many frames they have but rather must be parsed from beginning to end.
145pub trait FrameDataSource {
146    /// Get the width of the source images, in pixels.
147    fn width(&self) -> u32;
148    /// Get the height of the source images, in pixels.
149    fn height(&self) -> u32;
150    fn camera_name(&self) -> Option<&str> {
151        None
152    }
153    fn gamma(&self) -> Option<f32> {
154        None
155    }
156    /// Get the timestamp of the first frame.
157    ///
158    /// Note that (in case they can differ), this is the time
159    /// of the first frame rather than the creation time
160    /// in the metadata.
161    fn frame0_time(&self) -> Option<chrono::DateTime<chrono::FixedOffset>>;
162    /// Get the average framerate
163    ///
164    /// Value in frames per second.
165    fn average_framerate(&self) -> Option<f64>;
166    /// Set source to skip the first N frames.
167    ///
168    /// Note that this resets frame0_time accordingly.
169    fn skip_n_frames(&mut self, n_frames: usize) -> Result<()>;
170    /// Scan over the input images and estimate the luminance range
171    ///
172    /// Returns Ok<(min, max)> when successful.
173    fn estimate_luminance_range(&mut self) -> Result<(u16, u16)>;
174    /// Whether timestamps are available.
175    ///
176    /// If no timestamp is available, the frame "timestamp" with contain a
177    /// fraction of completeness.
178    fn has_timestamps(&self) -> bool;
179    /// A string describing the source of the timestamp data
180    fn timestamp_source(&self) -> &str;
181    /// Set when an SRT-file timestamp source ran out of usable stanzas
182    /// before covering every frame in the container, in which case the
183    /// source was truncated to the last complete group of pictures a stanza
184    /// is available for every frame of. `None` for sources not using an SRT
185    /// file, or when the SRT covered every frame.
186    fn srt_truncation(&self) -> Option<h264_source::SrtTruncation> {
187        None
188    }
189    /// Get an iterator over all frames in **decode order** (the order samples
190    /// are stored in the stream).
191    ///
192    /// For streams with B-frames this is *not* presentation order: the reported
193    /// per-frame [`Timestamp`]s are not monotonic in this order. Decode order is
194    /// what you want when feeding an H.264 decoder or re-muxing an H.264
195    /// bitstream (where reordering is reconstructed from per-sample composition
196    /// offsets). If you want frames in display order, use
197    /// [`Self::presentation_order_iter`].
198    fn decode_order_iter<'a>(&'a mut self) -> Box<dyn Iterator<Item = Result<FrameData>> + 'a>;
199
200    /// Get an iterator over all frames in **presentation (display) order**, with
201    /// monotonically non-decreasing timestamps.
202    ///
203    /// The default implementation is the identity over
204    /// [`Self::decode_order_iter`], which is correct for every source without
205    /// inter-frame reordering (FMF, TIFF, and H.264 without B-frames). Sources
206    /// that can reorder (H.264) override this to buffer and reorder frames.
207    ///
208    /// Returns an error up front when display order cannot be recovered (e.g. a
209    /// raw Annex B `.h264` stream carrying neither per-sample timestamps nor a
210    /// decodable picture order count), so callers fail loudly rather than
211    /// silently receiving decode order.
212    fn presentation_order_iter<'a>(
213        &'a mut self,
214    ) -> Result<Box<dyn Iterator<Item = Result<FrameData>> + 'a>> {
215        Ok(self.decode_order_iter())
216    }
217}
218
219/// A single frame of data, including `image` and `timestamp` fields.
220#[derive(Debug)]
221pub struct FrameData {
222    /// This is often called "PTS" (presentation time stamp).
223    timestamp: Timestamp,
224    image: ImageData,
225    buf_len: usize,
226    /// The number of the frame in the source, in **decode order**.
227    ///
228    /// Starts with 0. For streams with B-frames, decode order differs from
229    /// presentation (display) order, so this is *not* the display position; see
230    /// [`FrameData::poc`] and [`FrameDataSource::presentation_order_iter`].
231    idx: usize,
232    /// The frame's picture order count (POC) within its coded video sequence,
233    /// when known (H.264 sources only). This is the bitstream's own signal of
234    /// *display order*: within a coded video sequence, sorting frames by `poc`
235    /// yields presentation order. It is a relative rank, not a time value, and
236    /// resets at each IDR. `None` for sources without inter-frame reordering
237    /// (FMF, TIFF) or when it could not be reconstructed.
238    poc: Option<i64>,
239}
240
241#[derive(PartialEq, Debug, Clone, Copy)]
242pub enum Timestamp {
243    /// The timestamp, measured as the duration elapsed since the track onset
244    /// until the exposure started.
245    Duration(std::time::Duration),
246    /// In cases where no time is available, the fraction done.
247    Fraction(f32),
248}
249
250impl Timestamp {
251    pub fn unwrap_duration(&self) -> std::time::Duration {
252        match self {
253            Timestamp::Duration(d) => *d,
254            Timestamp::Fraction(_) => {
255                panic!("expected duration");
256            }
257        }
258    }
259}
260
261impl FrameData {
262    /// Get the timestamp, measured as the duration elapsed since the track onset
263    /// until the exposure started.
264    ///
265    /// This is often called "PTS" (presentation time stamp).
266    pub fn timestamp(&self) -> Timestamp {
267        self.timestamp
268    }
269    /// Get the image data
270    pub fn image(&self) -> &ImageData {
271        &self.image
272    }
273    /// Get the image data
274    pub fn into_image(self) -> ImageData {
275        self.image
276    }
277    /// Get the number of the bytes used in the source.
278    pub fn num_bytes(&self) -> usize {
279        self.buf_len
280    }
281    /// Get the number of the frame in the source, in decode order.
282    ///
283    /// Starts with 0. See the field docs on [`FrameData`] for the decode- vs
284    /// presentation-order distinction.
285    pub fn idx(&self) -> usize {
286        self.idx
287    }
288
289    /// Get the frame's picture order count (POC), when known.
290    ///
291    /// See the [`FrameData::poc`] field docs. `None` for non-H.264 sources or
292    /// when POC could not be reconstructed.
293    pub fn poc(&self) -> Option<i64> {
294        self.poc
295    }
296
297    pub fn decoded<'a>(&'a self) -> Option<DynamicFrame<'a>> {
298        match &self.image {
299            ImageData::Decoded(frame) => Some(frame.borrow()),
300            _ => None,
301        }
302    }
303
304    pub fn take_decoded(self) -> Option<DynamicFrameOwned> {
305        match self.image {
306            ImageData::Decoded(frame) => Some(frame),
307            _ => None,
308        }
309    }
310}
311
312/// The image data
313#[derive(Clone)]
314pub enum ImageData {
315    Decoded(DynamicFrameOwned),
316    Tiff(TiffImage),
317    EncodedH264(EncodedH264),
318}
319
320impl std::fmt::Debug for ImageData {
321    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
322        match self {
323            ImageData::Decoded(_) => {
324                write!(f, "ImageData::Decoded")
325            }
326            ImageData::Tiff(_) => {
327                write!(f, "ImageData::Tiff")
328            }
329            ImageData::EncodedH264(_) => {
330                write!(f, "ImageData::EncodedH264")
331            }
332        }
333    }
334}
335
336#[derive(Clone, PartialEq)]
337pub enum H264EncodingVariant {
338    /// single large buffer with Annex B headers
339    AnnexB(Vec<u8>),
340    /// single large buffer with AVCC headers
341    Avcc(Vec<u8>),
342    /// multiple buffers with just NAL unit data
343    RawEbsp(Vec<Vec<u8>>),
344}
345
346impl std::fmt::Debug for H264EncodingVariant {
347    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
348        match self {
349            Self::AnnexB(buf) => write!(f, "H264EncodingVariant::AnnexB({} bytes)", buf.len()),
350            Self::Avcc(buf) => write!(f, "H264EncodingVariant::Avcc({} bytes)", buf.len()),
351            Self::RawEbsp(bufs) => {
352                write!(f, "H264EncodingVariant::RawEbsp({} buffers)", bufs.len())
353            }
354        }
355    }
356}
357
358#[derive(Clone, PartialEq, Debug)]
359pub struct EncodedH264 {
360    pub data: H264EncodingVariant,
361    pub has_precision_timestamp: bool,
362}
363
364#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
365pub enum TimestampSource {
366    #[default]
367    BestGuess,
368    /// H264 contains FrameInfo (`strawlab.org/89H`) supplemental enhancement
369    /// information NAL units and this source uses the receive time of the
370    /// receiving computer. (Note the RTP time of the sender is unused.)
371    FrameInfoRecvTime,
372    /// H264 contains FrameInfo (`strawlab.org/89H`) supplemental enhancement
373    /// information NAL units and this source uses the RTP time of the
374    /// sending camera with best-guess offset to receiving computer.
375    FrameInfoRtp,
376    /// Use the Presentation Time Stamp (PTS) of the MP4 data.
377    Mp4Pts,
378    /// H264 contains "MISPmicrosectime" supplemental enhancement information
379    /// NAL units and this source uses it.
380    MispMicrosectime,
381    /// Using timing from a specific schema of .srt files.
382    SrtFile,
383    /// Simply multiply frame number by average framerate
384    FixedFramerate,
385}
386
387trait MyAsStr {
388    fn as_str(&self) -> &'static str;
389}
390
391impl MyAsStr for Option<TimestampSource> {
392    fn as_str(&self) -> &'static str {
393        use TimestampSource::*;
394        match self {
395            Some(BestGuess) => "(best guess)",
396            Some(FrameInfoRecvTime) => "FrameInfo receive time",
397            Some(FrameInfoRtp) => "FrameInfo RTP",
398            Some(Mp4Pts) => "MP4 PTS",
399            Some(MispMicrosectime) => "MISPmicrosectime",
400            Some(SrtFile) => "SRT file",
401            Some(FixedFramerate) => "frame number multiplied by average frame rate",
402            None => "(no timestamps)",
403        }
404    }
405}
406
407/// Builder for frame sources. Use this to set various options on the frame
408/// source.
409pub struct FrameSourceBuilder {
410    input: PathBuf,
411    do_decode_h264: bool,
412    timestamp_source: TimestampSource,
413    srt_file_path: Option<PathBuf>,
414    show_progress: bool,
415}
416
417impl FrameSourceBuilder {
418    pub fn new<P: AsRef<std::path::Path>>(input: P) -> Self {
419        Self {
420            input: PathBuf::from(input.as_ref()),
421            do_decode_h264: true,
422            timestamp_source: TimestampSource::BestGuess,
423            srt_file_path: None,
424            show_progress: false,
425        }
426    }
427    pub fn do_decode_h264(self, do_decode_h264: bool) -> Self {
428        Self {
429            do_decode_h264,
430            ..self
431        }
432    }
433    pub fn timestamp_source(self, timestamp_source: TimestampSource) -> Self {
434        Self {
435            timestamp_source,
436            ..self
437        }
438    }
439    pub fn srt_file_path(self, srt_file_path: Option<PathBuf>) -> Self {
440        Self {
441            srt_file_path,
442            ..self
443        }
444    }
445    pub fn show_progress(self, show_progress: bool) -> Self {
446        Self {
447            show_progress,
448            ..self
449        }
450    }
451    /// Create a [FrameDataSource]
452    pub fn build_source(self) -> Result<Box<dyn FrameDataSource>> {
453        build_frame_source(
454            self.input,
455            self.do_decode_h264,
456            self.timestamp_source,
457            self.srt_file_path,
458            self.show_progress,
459        )
460    }
461    pub fn build_h264_in_mp4_source(
462        self,
463    ) -> Result<h264_source::H264Source<mp4_source::Mp4Source>> {
464        mp4_source::open_h264_in_mp4(
465            self.input,
466            self.do_decode_h264,
467            self.timestamp_source,
468            self.srt_file_path,
469            self.show_progress,
470            None,
471        )
472    }
473    pub fn build_h264_in_mp4_source_with_preparser(
474        self,
475        preparser: Box<dyn h264_source::H264Preparser>,
476    ) -> Result<h264_source::H264Source<mp4_source::Mp4Source>> {
477        mp4_source::open_h264_in_mp4(
478            self.input,
479            self.do_decode_h264,
480            self.timestamp_source,
481            self.srt_file_path,
482            self.show_progress,
483            Some(preparser),
484        )
485    }
486    /// Build a source for a raw "Annex B" `.h264` file.
487    pub fn build_h264_annexb_source(
488        self,
489    ) -> Result<h264_source::H264Source<h264_source::H264AnnexBSource>> {
490        h264_source::from_annexb_path_with_timestamp_source(
491            self.input,
492            self.do_decode_h264,
493            self.timestamp_source,
494            self.srt_file_path,
495            self.show_progress,
496        )
497    }
498    pub fn build_mkv_source(
499        self,
500    ) -> Result<strand_cam_mkv_source::StrandCamMkvSource<std::io::BufReader<std::fs::File>>> {
501        if self.srt_file_path.is_some() {
502            return Err(Error::NoSrtSupportForFileType);
503        }
504        if self.show_progress {
505            return Err(Error::UnsupportedOption);
506        }
507        strand_cam_mkv_source::mkv_source_from_path_with_timestamp_source(
508            self.input,
509            self.do_decode_h264,
510            self.timestamp_source,
511        )
512    }
513}
514
515fn build_frame_source(
516    input_path: PathBuf,
517    do_decode_h264: bool,
518    timestamp_source: TimestampSource,
519    srt_file_path: Option<PathBuf>,
520    show_progress: bool,
521) -> Result<Box<dyn FrameDataSource>> {
522    let is_file = std::fs::metadata(&input_path)?.is_file();
523    if is_file {
524        if let Some(extension) = input_path.extension() {
525            let lower_ext = extension.to_str().map(|x| x.to_string().to_lowercase());
526            match lower_ext.as_deref() {
527                Some("mkv") => {
528                    if srt_file_path.is_some() {
529                        return Err(Error::NoSrtSupportForFileType);
530                    }
531                    if show_progress {
532                        return Err(Error::UnsupportedOption);
533                    }
534                    let mkv_video =
535                        strand_cam_mkv_source::mkv_source_from_path_with_timestamp_source(
536                            &input_path,
537                            do_decode_h264,
538                            timestamp_source,
539                        )?;
540                    return Ok(Box::new(mkv_video));
541                }
542                Some("mp4") => {
543                    let mp4_video = mp4_source::open_h264_in_mp4(
544                        &input_path,
545                        do_decode_h264,
546                        timestamp_source,
547                        srt_file_path,
548                        show_progress,
549                        None,
550                    )?;
551                    return Ok(Box::new(mp4_video));
552                }
553                Some("h264") => {
554                    if srt_file_path.is_some() {
555                        return Err(Error::NoSrtSupportForFileType);
556                    }
557                    let h264_video = h264_source::from_annexb_path_with_timestamp_source(
558                        &input_path,
559                        do_decode_h264,
560                        timestamp_source,
561                        None,
562                        show_progress,
563                    )?;
564                    return Ok(Box::new(h264_video));
565                }
566                _ => {}
567            }
568        }
569        let fname_lower = input_path.to_string_lossy().to_lowercase();
570        if fname_lower.ends_with(".fmf") || fname_lower.ends_with(".fmf.gz") {
571            let fmf_video = fmf_source::from_path(&input_path)?;
572            return Ok(Box::new(fmf_video));
573        }
574        Err(Error::UnknownExtensionForFile(input_path))
575    } else {
576        let dirname = input_path;
577
578        if !std::fs::metadata(&dirname)?.is_dir() {
579            return Err(Error::TiffStackNotDir(dirname));
580        }
581        let pattern = dirname.join("*.tif");
582        if srt_file_path.is_some() {
583            return Err(Error::NoSrtSupportForFileType);
584        }
585        let stack = pv_tiff_stack::from_path_pattern(pattern.to_str().unwrap())?;
586        Ok(Box::new(stack))
587    }
588}