Skip to main content

rivet/
probe.rs

1//! Inspect an input without transcoding it.
2//!
3//! Demuxes just the container header + audio track metadata and reports the
4//! video codec, dimensions, frame rate, pixel format, and audio stream
5//! shape. Works across every container the [`container`] crate supports
6//! (MP4/MOV, MKV/WebM, AVI, MPEG-TS).
7
8use std::path::Path;
9
10use anyhow::{Context, Result};
11
12use container::streaming;
13
14/// Probed media metadata.
15#[derive(Debug, Clone)]
16pub struct MediaInfo {
17    /// Detected container label: `"mp4"`, `"mkv"`, `"avi"`, or `"ts"`.
18    pub container: String,
19    /// Lower-cased video codec label (e.g. `"h264"`, `"hevc"`, `"av1"`).
20    pub video_codec: String,
21    /// Video width in pixels (0 if the container did not record it).
22    pub width: u32,
23    /// Video height in pixels (0 if the container did not record it).
24    pub height: u32,
25    /// Frame rate in frames per second.
26    pub frame_rate: f64,
27    /// Duration in seconds (0.0 if the container did not record it).
28    pub duration: f64,
29    /// Pixel format, e.g. `"Yuv420p"` / `"Yuv420p10le"`.
30    pub pixel_format: String,
31    /// Audio stream metadata, if present.
32    pub audio: Option<AudioStreamInfo>,
33}
34
35/// Audio stream metadata.
36#[derive(Debug, Clone)]
37pub struct AudioStreamInfo {
38    /// Lower-cased audio codec label (e.g. `"aac"`, `"opus"`, `"mp3"`).
39    pub codec: String,
40    /// Sample rate in Hz.
41    pub sample_rate: u32,
42    /// Channel count.
43    pub channels: u16,
44}
45
46/// Probe an input file.
47pub fn probe_file(input: impl AsRef<Path>) -> Result<MediaInfo> {
48    let input = input.as_ref();
49    let bytes = std::fs::read(input)
50        .with_context(|| format!("reading input file {}", input.display()))?;
51    probe_bytes(&bytes)
52}
53
54/// Probe an in-memory input buffer.
55pub fn probe_bytes(input: &[u8]) -> Result<MediaInfo> {
56    let demuxer = streaming::demux_streaming(input).context("demux")?;
57    let header = demuxer.header();
58
59    let audio = demuxer.audio().map(|t| AudioStreamInfo {
60        codec: t.codec.to_ascii_lowercase(),
61        sample_rate: t.sample_rate,
62        channels: t.channels,
63    });
64
65    Ok(MediaInfo {
66        container: detect_container(input).to_string(),
67        video_codec: header.codec.to_ascii_lowercase(),
68        width: header.info.width,
69        height: header.info.height,
70        frame_rate: header.info.frame_rate,
71        duration: header.info.duration,
72        pixel_format: format!("{:?}", header.info.pixel_format),
73        audio,
74    })
75}
76
77/// Magic-byte container detector — mirrors the dispatch in
78/// [`container::streaming::demux_streaming`] so the reported label matches
79/// the demuxer that was actually used.
80fn detect_container(data: &[u8]) -> &'static str {
81    if data.len() < 12 {
82        return "unknown";
83    }
84    if &data[4..8] == b"ftyp" || &data[4..8] == b"moov" || &data[4..8] == b"mdat" {
85        return "mp4";
86    }
87    if data[0] == 0x1A && data[1] == 0x45 && data[2] == 0xDF && data[3] == 0xA3 {
88        return "mkv";
89    }
90    if &data[..4] == b"RIFF" && &data[8..12] == b"AVI " {
91        return "avi";
92    }
93    if data[0] == 0x47
94        && data.len() > 188
95        && data[188] == 0x47
96        && (data.len() <= 376 || data[376] == 0x47)
97    {
98        return "ts";
99    }
100    "unknown"
101}