codec/
probe.rs

1//! Stream probe — analyze video files without FFmpeg.
2//!
3//! Extracts codec, resolution, frame rate, duration, and container
4//! metadata from MP4/MOV files using the mp4 crate.
5
6use anyhow::{Context, Result};
7use std::io::Cursor;
8
9use crate::frame::{
10    ColorMetadata, ColorSpace, ContentLightLevel, MasteringDisplay, PixelFormat, StreamInfo,
11};
12use crate::hevc_sei;
13
14#[derive(Debug, Clone)]
15pub struct ProbeResult {
16    pub stream_info: StreamInfo,
17    pub container: String,
18    pub audio_codec: Option<String>,
19    pub audio_sample_rate: Option<u32>,
20    pub audio_channels: Option<u16>,
21    pub file_size: u64,
22    pub metadata: std::collections::HashMap<String, String>,
23}
24
25pub fn probe_mp4(data: &[u8]) -> Result<ProbeResult> {
26    let size = data.len() as u64;
27    let cursor = Cursor::new(data);
28    let reader =
29        mp4::Mp4Reader::read_header(cursor, size).context("reading MP4 header for probe")?;
30
31    let video_track = reader
32        .tracks()
33        .values()
34        .find(|t| t.track_type().ok() == Some(mp4::TrackType::Video))
35        .context("no video track")?;
36
37    let codec = match video_track.media_type() {
38        Ok(mp4::MediaType::H264) => "h264",
39        Ok(mp4::MediaType::H265) => "h265",
40        Ok(mp4::MediaType::VP9) => "vp9",
41        _ => "unknown",
42    };
43
44    let width = video_track.width() as u32;
45    let height = video_track.height() as u32;
46    let sample_count = video_track.sample_count();
47    let duration = video_track.duration().as_secs_f64();
48    let frame_rate = if duration > 0.0 {
49        sample_count as f64 / duration
50    } else {
51        30.0
52    };
53    let bitrate = video_track.bitrate() as u64;
54
55    let audio_track = reader
56        .tracks()
57        .values()
58        .find(|t| t.track_type().ok() == Some(mp4::TrackType::Audio));
59
60    let audio_codec = audio_track.and_then(|t| t.media_type().ok().map(|mt| format!("{mt:?}")));
61    let audio_sample_rate: Option<u32> = None;
62    let audio_channels: Option<u16> = None;
63
64    // Squad-21: extract HDR static-metadata boxes (`mdcv`, `clli`) from
65    // the visual sample entry. These are the canonical container-side
66    // HDR10 carriers — without surfacing them, the muxer can't write
67    // them on the output and Apple devices fall back to BT.709 limited.
68    let probe_color = probe_mp4_visual_color_metadata(data);
69    let color_metadata = ColorMetadata {
70        mastering_display: probe_color.mastering_display,
71        content_light_level: probe_color.content_light_level,
72        ..ColorMetadata::default()
73    };
74
75    let stream_info = StreamInfo {
76        codec: codec.to_string(),
77        width,
78        height,
79        frame_rate,
80        duration,
81        pixel_format: PixelFormat::Yuv420p,
82        color_space: ColorSpace::Bt709,
83        total_frames: sample_count as u64,
84        bitrate,
85        color_metadata,
86    };
87
88    Ok(ProbeResult {
89        stream_info,
90        container: "mp4".to_string(),
91        audio_codec,
92        audio_sample_rate,
93        audio_channels,
94        file_size: size,
95        metadata: std::collections::HashMap::new(),
96    })
97}
98
99/// Squad-21: HDR static metadata pulled from MP4 visual sample-entry
100/// boxes (`mdcv`, `clli`). Returns `None` for SDR sources or non-MP4
101/// inputs.
102#[derive(Debug, Default, Clone, Copy)]
103struct ProbeMp4VisualColorMetadata {
104    mastering_display: Option<MasteringDisplay>,
105    content_light_level: Option<ContentLightLevel>,
106}
107
108fn probe_mp4_visual_color_metadata(data: &[u8]) -> ProbeMp4VisualColorMetadata {
109    let path: &[&[u8; 4]] = &[b"moov", b"trak", b"mdia", b"minf", b"stbl", b"stsd"];
110    let Some(stsd_body) = find_box_body(data, path) else {
111        return ProbeMp4VisualColorMetadata::default();
112    };
113    if stsd_body.len() < 16 {
114        return ProbeMp4VisualColorMetadata::default();
115    }
116
117    let mut pos = 8;
118    while pos + 8 <= stsd_body.len() {
119        let entry_size = u32::from_be_bytes([
120            stsd_body[pos],
121            stsd_body[pos + 1],
122            stsd_body[pos + 2],
123            stsd_body[pos + 3],
124        ]) as usize;
125        if entry_size < 8 || pos.saturating_add(entry_size) > stsd_body.len() {
126            break;
127        }
128        let entry_type: [u8; 4] = match stsd_body[pos + 4..pos + 8].try_into() {
129            Ok(v) => v,
130            Err(_) => break,
131        };
132        let is_visual = matches!(
133            &entry_type,
134            b"av01"
135                | b"avc1"
136                | b"avc3"
137                | b"hvc1"
138                | b"hev1"
139                | b"hvc2"
140                | b"hev2"
141                | b"dvh1"
142                | b"dvhe"
143                | b"vp08"
144                | b"vp09"
145                | b"apcn"
146                | b"apch"
147                | b"apcs"
148                | b"apco"
149                | b"ap4h"
150                | b"ap4x"
151        );
152        if !is_visual {
153            pos = pos.saturating_add(entry_size);
154            continue;
155        }
156        let end = pos.saturating_add(entry_size);
157        let child_start = pos + 8 + 78;
158        if child_start >= end {
159            return ProbeMp4VisualColorMetadata::default();
160        }
161        let children = &stsd_body[child_start..end];
162        let mut out = ProbeMp4VisualColorMetadata::default();
163        if let Some(mdcv) = find_direct_child(children, b"mdcv")
164            && mdcv.len() >= 24
165        {
166            let u16be = |o: usize| u16::from_be_bytes([mdcv[o], mdcv[o + 1]]);
167            let u32be =
168                |o: usize| u32::from_be_bytes([mdcv[o], mdcv[o + 1], mdcv[o + 2], mdcv[o + 3]]);
169            out.mastering_display = Some(MasteringDisplay {
170                primaries_g_x: u16be(0),
171                primaries_g_y: u16be(2),
172                primaries_b_x: u16be(4),
173                primaries_b_y: u16be(6),
174                primaries_r_x: u16be(8),
175                primaries_r_y: u16be(10),
176                white_point_x: u16be(12),
177                white_point_y: u16be(14),
178                max_luminance: u32be(16),
179                min_luminance: u32be(20),
180            });
181        }
182        if let Some(clli) = find_direct_child(children, b"clli")
183            && clli.len() >= 4
184        {
185            out.content_light_level = Some(ContentLightLevel {
186                max_cll: u16::from_be_bytes([clli[0], clli[1]]),
187                max_fall: u16::from_be_bytes([clli[2], clli[3]]),
188            });
189        }
190        return out;
191    }
192    ProbeMp4VisualColorMetadata::default()
193}
194
195/// Walk the ISOBMFF box tree following `path` and return the body bytes
196/// of the deepest box (or None if any hop is missing). Local helper —
197/// duplicates `container::demux::find_box_body` to avoid making
198/// `codec` depend on `container`.
199fn find_box_body<'a>(data: &'a [u8], path: &[&[u8; 4]]) -> Option<&'a [u8]> {
200    let mut slice = data;
201    for (i, target) in path.iter().enumerate() {
202        let found = find_direct_child(slice, target)?;
203        if i + 1 == path.len() {
204            return Some(found);
205        }
206        slice = found;
207    }
208    None
209}
210
211fn find_direct_child<'a>(data: &'a [u8], target: &[u8; 4]) -> Option<&'a [u8]> {
212    let mut pos = 0;
213    while pos + 8 <= data.len() {
214        let size =
215            u32::from_be_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]) as usize;
216        let btype = &data[pos + 4..pos + 8];
217        if size < 8 || pos.checked_add(size).is_none_or(|end| end > data.len()) {
218            return None;
219        }
220        if btype == target {
221            return Some(&data[pos + 8..pos + size]);
222        }
223        pos += size;
224    }
225    None
226}
227
228/// Re-export of the SEI parser entry-point for downstream callers
229/// that want to scan an Annex-B HEVC bitstream directly without
230/// constructing a decoder. The decoder also folds this internally
231/// (see `decode::hevc_de265::De265Decoder::new`).
232pub use hevc_sei::parse_annexb as parse_hevc_hdr_sei;
233
234pub fn detect_container(data: &[u8]) -> &'static str {
235    if data.len() < 12 {
236        return "unknown";
237    }
238    // MP4/MOV: ftyp box at offset 4
239    if &data[4..8] == b"ftyp" {
240        return "mp4";
241    }
242    // MKV/WebM: EBML header
243    if data[0] == 0x1A && data[1] == 0x45 && data[2] == 0xDF && data[3] == 0xA3 {
244        return "mkv";
245    }
246    // AVI: RIFF header
247    if &data[0..4] == b"RIFF" && data.len() > 11 && &data[8..12] == b"AVI " {
248        return "avi";
249    }
250    "unknown"
251}
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256
257    #[test]
258    fn test_detect_container_mp4() {
259        let mut data = vec![0u8; 16];
260        data[4..8].copy_from_slice(b"ftyp");
261        assert_eq!(detect_container(&data), "mp4");
262    }
263
264    #[test]
265    fn test_detect_container_mkv() {
266        let data = vec![
267            0x1A, 0x45, 0xDF, 0xA3, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
268        ];
269        assert_eq!(detect_container(&data), "mkv");
270    }
271
272    #[test]
273    fn test_detect_container_avi() {
274        let mut data = vec![0u8; 16];
275        data[0..4].copy_from_slice(b"RIFF");
276        data[8..12].copy_from_slice(b"AVI ");
277        assert_eq!(detect_container(&data), "avi");
278    }
279
280    #[test]
281    fn test_detect_container_unknown() {
282        let data = vec![0xFF; 16];
283        assert_eq!(detect_container(&data), "unknown");
284    }
285
286    #[test]
287    fn test_detect_container_short() {
288        let data = vec![0u8; 4];
289        assert_eq!(detect_container(&data), "unknown");
290    }
291}
codec/probe.rs

codec/
probe.rs
