Skip to main content

viser_ffmpeg/
probe.rs

1use serde::{Deserialize, Serialize};
2use tokio::process::Command;
3
4use crate::ffprobe_path;
5
6/// Parsed result of an `ffprobe` run: container format plus all streams.
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct ProbeResult {
9    /// Container-level format information.
10    pub format: FormatInfo,
11    /// All streams (video, audio, subtitle) found in the file.
12    pub streams: Vec<StreamInfo>,
13}
14
15/// Container-level metadata reported by ffprobe.
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct FormatInfo {
18    /// Probed file path.
19    pub filename: String,
20    /// Short container format name (e.g. `"mov,mp4,m4a,..."`).
21    pub format_name: String,
22    /// Human-readable container format name.
23    pub format_long_name: String,
24    /// Total duration in seconds.
25    pub duration: f64, // seconds
26    /// File size in bytes.
27    pub size: i64, // bytes
28    /// Overall bitrate in bits per second.
29    pub bit_rate: i64, // bits/sec
30    /// ffprobe's confidence score for the detected format.
31    pub probe_score: i32,
32}
33
34/// Per-stream metadata reported by ffprobe.
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct StreamInfo {
37    /// Stream index within the container.
38    pub index: i32,
39    /// Short codec name (e.g. `"h264"`).
40    pub codec_name: String,
41    /// Human-readable codec name.
42    pub codec_long_name: String,
43    /// Stream type: `"video"`, `"audio"`, or `"subtitle"`.
44    pub codec_type: String, // "video", "audio", "subtitle"
45    /// Codec profile (e.g. `"High"`).
46    pub profile: String,
47    /// Pixel width (video).
48    pub width: i32,
49    /// Pixel height (video).
50    pub height: i32,
51    /// Pixel format (e.g. `"yuv420p"`).
52    pub pix_fmt: String,
53    /// Codec level.
54    pub level: i32,
55    /// Field/scan order (e.g. `"progressive"`).
56    pub field_order: String,
57    /// Color range (e.g. `"tv"`/`"pc"`).
58    pub color_range: String,
59    /// Color matrix / space.
60    pub color_space: String,
61    /// Color transfer characteristics (e.g. `"smpte2084"` for PQ).
62    pub color_transfer: String,
63    /// Color primaries (e.g. `"bt2020"`).
64    pub color_primaries: String,
65    /// Stream duration in seconds.
66    pub duration: f64, // seconds
67    /// Stream bitrate in bits per second.
68    pub bit_rate: i64, // bits/sec
69    /// Number of frames, when known.
70    pub nb_frames: i32,
71    /// Raw frame rate as a rational string (e.g. `"25/1"`).
72    pub r_frame_rate: String, // e.g. "25/1"
73    /// Average frame rate as a rational string (e.g. `"25/1"`).
74    pub avg_frame_rate: String, // e.g. "25/1"
75    /// Audio sample rate in Hz.
76    pub sample_rate: i32, // audio
77    /// Audio channel count.
78    pub channels: i32, // audio
79    /// Audio channel layout (e.g. `"stereo"`).
80    pub channel_layout: String, // audio
81    /// Bits per raw sample (used to detect high-bit-depth/HDR content).
82    pub bits_per_raw_sample: i32,
83}
84
85impl StreamInfo {
86    /// Returns an error unless this is a video stream with positive dimensions.
87    pub fn validate(&self) -> anyhow::Result<()> {
88        if self.codec_type != "video" {
89            anyhow::bail!("not a video stream (type={})", self.codec_type);
90        }
91        if self.width <= 0 || self.height <= 0 {
92            anyhow::bail!("invalid dimensions: {}x{}", self.width, self.height);
93        }
94        Ok(())
95    }
96
97    /// Frame rate parsed from r_frame_rate (e.g. "50/1" -> 50.0).
98    pub fn fps(&self) -> f64 {
99        parse_rational(&self.r_frame_rate)
100    }
101
102    /// Returns `"WIDTHxHEIGHT"`, or an empty string if dimensions are unknown.
103    pub fn resolution_str(&self) -> String {
104        if self.width == 0 || self.height == 0 {
105            String::new()
106        } else {
107            format!("{}x{}", self.width, self.height)
108        }
109    }
110
111    /// Detects HDR type from color metadata: `"PQ"`, `"HLG"`, `"BT.2020"`, or `None` for SDR.
112    pub fn hdr_kind(&self) -> Option<&'static str> {
113        let transfer = self.color_transfer.to_ascii_lowercase();
114        if transfer == "smpte2084" {
115            return Some("PQ");
116        }
117        if transfer == "arib-std-b67" {
118            return Some("HLG");
119        }
120
121        let primaries_bt2020 = self.color_primaries.eq_ignore_ascii_case("bt2020");
122        let high_bit_depth = self.bits_per_raw_sample >= 10
123            || self.pix_fmt.contains("10")
124            || self.pix_fmt.contains("12")
125            || self.pix_fmt.contains("16");
126        if primaries_bt2020 && high_bit_depth {
127            return Some("BT.2020");
128        }
129
130        None
131    }
132
133    /// Returns `true` if the stream carries HDR color metadata.
134    pub fn is_hdr(&self) -> bool {
135        self.hdr_kind().is_some()
136    }
137}
138
139impl FormatInfo {
140    /// Container duration as a `Duration`.
141    pub fn duration_secs(&self) -> std::time::Duration {
142        std::time::Duration::from_secs_f64(self.duration)
143    }
144}
145
146impl ProbeResult {
147    /// First video stream, if any.
148    pub fn video_stream(&self) -> Option<&StreamInfo> {
149        self.streams.iter().find(|s| s.codec_type == "video")
150    }
151
152    /// First audio stream, if any.
153    pub fn audio_stream(&self) -> Option<&StreamInfo> {
154        self.streams.iter().find(|s| s.codec_type == "audio")
155    }
156}
157
158/// Runs ffprobe on the given file and returns parsed results.
159pub async fn probe(path: &str) -> anyhow::Result<ProbeResult> {
160    let args = ["-v", "error", "-print_format", "json", "-show_format", "-show_streams", path];
161
162    let output = Command::new(ffprobe_path())
163        .args(args)
164        .stderr(std::process::Stdio::piped())
165        .output()
166        .await?;
167
168    if !output.status.success() {
169        let stderr = String::from_utf8_lossy(&output.stderr);
170        anyhow::bail!("ffprobe failed for {path}: {stderr}");
171    }
172
173    let raw: ProbeJsonRaw = serde_json::from_slice(&output.stdout)
174        .map_err(|e| anyhow::anyhow!("failed to parse ffprobe output: {e}"))?;
175
176    Ok(convert_probe(raw))
177}
178
179// Raw ffprobe JSON — numbers come as strings
180#[derive(Deserialize)]
181struct ProbeJsonRaw {
182    format: ProbeFormatRaw,
183    streams: Vec<ProbeStreamRaw>,
184}
185
186#[derive(Deserialize)]
187struct ProbeFormatRaw {
188    #[serde(default)]
189    filename: String,
190    #[serde(default)]
191    format_name: String,
192    #[serde(default)]
193    format_long_name: String,
194    #[serde(default)]
195    duration: String,
196    #[serde(default)]
197    size: String,
198    #[serde(default)]
199    bit_rate: String,
200    #[serde(default)]
201    probe_score: i32,
202}
203
204#[derive(Deserialize)]
205struct ProbeStreamRaw {
206    #[serde(default)]
207    index: i32,
208    #[serde(default)]
209    codec_name: String,
210    #[serde(default)]
211    codec_long_name: String,
212    #[serde(default)]
213    codec_type: String,
214    #[serde(default)]
215    profile: String,
216    #[serde(default)]
217    width: i32,
218    #[serde(default)]
219    height: i32,
220    #[serde(default)]
221    pix_fmt: String,
222    #[serde(default)]
223    level: i32,
224    #[serde(default)]
225    field_order: String,
226    #[serde(default)]
227    color_range: String,
228    #[serde(default)]
229    color_space: String,
230    #[serde(default)]
231    color_transfer: String,
232    #[serde(default)]
233    color_primaries: String,
234    #[serde(default)]
235    duration: String,
236    #[serde(default)]
237    bit_rate: String,
238    #[serde(default)]
239    nb_frames: String,
240    #[serde(default)]
241    r_frame_rate: String,
242    #[serde(default)]
243    avg_frame_rate: String,
244    #[serde(default)]
245    sample_rate: String,
246    #[serde(default)]
247    channels: i32,
248    #[serde(default)]
249    channel_layout: String,
250    #[serde(default)]
251    bits_per_raw_sample: String,
252}
253
254fn convert_probe(raw: ProbeJsonRaw) -> ProbeResult {
255    let format = FormatInfo {
256        filename: raw.format.filename,
257        format_name: raw.format.format_name,
258        format_long_name: raw.format.format_long_name,
259        duration: raw.format.duration.parse().unwrap_or(0.0),
260        size: raw.format.size.parse().unwrap_or(0),
261        bit_rate: raw.format.bit_rate.parse().unwrap_or(0),
262        probe_score: raw.format.probe_score,
263    };
264
265    let streams = raw
266        .streams
267        .into_iter()
268        .map(|s| StreamInfo {
269            index: s.index,
270            codec_name: s.codec_name,
271            codec_long_name: s.codec_long_name,
272            codec_type: s.codec_type,
273            profile: s.profile,
274            width: s.width,
275            height: s.height,
276            pix_fmt: s.pix_fmt,
277            level: s.level,
278            field_order: s.field_order,
279            color_range: s.color_range,
280            color_space: s.color_space,
281            color_transfer: s.color_transfer,
282            color_primaries: s.color_primaries,
283            duration: s.duration.parse().unwrap_or(0.0),
284            bit_rate: s.bit_rate.parse().unwrap_or(0),
285            nb_frames: s.nb_frames.parse().unwrap_or(0),
286            r_frame_rate: s.r_frame_rate,
287            avg_frame_rate: s.avg_frame_rate,
288            sample_rate: s.sample_rate.parse().unwrap_or(0),
289            channels: s.channels,
290            channel_layout: s.channel_layout,
291            bits_per_raw_sample: s.bits_per_raw_sample.parse().unwrap_or(0),
292        })
293        .collect();
294
295    ProbeResult { format, streams }
296}
297
298fn parse_rational(s: &str) -> f64 {
299    if let Some((num_s, den_s)) = s.split_once('/') {
300        let num: f64 = num_s.parse().unwrap_or(0.0);
301        let den: f64 = den_s.parse().unwrap_or(0.0);
302        if den != 0.0 { num / den } else { 0.0 }
303    } else {
304        s.parse().unwrap_or(0.0)
305    }
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311
312    fn video_stream() -> StreamInfo {
313        StreamInfo {
314            index: 0,
315            codec_name: "h264".into(),
316            codec_long_name: String::new(),
317            codec_type: "video".into(),
318            profile: String::new(),
319            width: 1920,
320            height: 1080,
321            pix_fmt: "yuv420p".into(),
322            level: 0,
323            field_order: String::new(),
324            color_range: String::new(),
325            color_space: String::new(),
326            color_transfer: String::new(),
327            color_primaries: String::new(),
328            duration: 0.0,
329            bit_rate: 0,
330            nb_frames: 0,
331            r_frame_rate: "24/1".into(),
332            avg_frame_rate: "24/1".into(),
333            sample_rate: 0,
334            channels: 0,
335            channel_layout: String::new(),
336            bits_per_raw_sample: 8,
337        }
338    }
339
340    #[test]
341    fn test_hdr_kind_detects_pq() {
342        let mut stream = video_stream();
343        stream.color_transfer = "smpte2084".into();
344        assert_eq!(stream.hdr_kind(), Some("PQ"));
345        assert!(stream.is_hdr());
346    }
347
348    #[test]
349    fn test_hdr_kind_detects_hlg() {
350        let mut stream = video_stream();
351        stream.color_transfer = "arib-std-b67".into();
352        assert_eq!(stream.hdr_kind(), Some("HLG"));
353    }
354
355    #[test]
356    fn test_hdr_kind_detects_bt2020_high_bit_depth() {
357        let mut stream = video_stream();
358        stream.color_primaries = "bt2020".into();
359        stream.pix_fmt = "yuv420p10le".into();
360        assert_eq!(stream.hdr_kind(), Some("BT.2020"));
361    }
362
363    #[test]
364    fn test_hdr_kind_ignores_sdr() {
365        assert_eq!(video_stream().hdr_kind(), None);
366    }
367}