Skip to main content

viser_ffmpeg/
probe.rs

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