frame_source/
h264_poc.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! Reconstruct the H.264 **picture order count** (POC) from a bitstream.
5//!
6//! The POC (ITU-T H.264 §8.2.1) is the one signal that cannot lie about a
7//! stream's true *display order*: every slice header carries enough information
8//! to recover the relative presentation order of samples, independent of any
9//! container metadata (`stts`/`ctts`) or of what a (possibly buggy) writer put
10//! in a per-frame timestamp SEI. Within a coded video sequence (delimited by
11//! IDR pictures), sorting frames by POC yields presentation order.
12//!
13//! This module is used both to reorder frames into presentation order (see
14//! [`crate::FrameDataSource::presentation_order_iter`]) and by
15//! `mp4-bframe-doctor` to detect/repair files whose timing disagrees with the
16//! bitstream's real display order.
17
18use h264_reader::{
19    Context as H264ParsingContext,
20    nal::{
21        Nal, RefNal, UnitType,
22        pps::PicParameterSet,
23        slice::{PicOrderCountLsb, SliceHeader},
24        sps::{PicOrderCntType, SeqParameterSet},
25    },
26};
27
28use crate::{
29    Error, Result,
30    h264_source::{H264Source, SeekableH264Source},
31};
32
33/// Reconstructs picture order count (POC) for `pic_order_cnt_type == 0`
34/// streams (ITU-T H.264 §8.2.1.1), which covers essentially all cameras and
35/// software/hardware H.264 encoders in practice.
36pub(crate) struct PocDecoder {
37    max_poc_lsb: i64,
38    prev_poc_msb: i64,
39    prev_poc_lsb: i64,
40}
41
42impl PocDecoder {
43    fn new(log2_max_pic_order_cnt_lsb_minus4: u8) -> Self {
44        Self {
45            max_poc_lsb: 1i64 << (log2_max_pic_order_cnt_lsb_minus4 as i64 + 4),
46            prev_poc_msb: 0,
47            prev_poc_lsb: 0,
48        }
49    }
50
51    /// Feed the next sample's slice header info, in decode order, and get
52    /// back its POC.
53    fn next_poc(&mut self, is_idr: bool, nal_ref_idc: u8, poc_lsb: i64) -> i64 {
54        if is_idr {
55            self.prev_poc_msb = 0;
56            self.prev_poc_lsb = 0;
57        }
58
59        let half_max = self.max_poc_lsb / 2;
60        let poc_msb = if poc_lsb < self.prev_poc_lsb && (self.prev_poc_lsb - poc_lsb) >= half_max {
61            self.prev_poc_msb + self.max_poc_lsb
62        } else if poc_lsb > self.prev_poc_lsb && (poc_lsb - self.prev_poc_lsb) > half_max {
63            self.prev_poc_msb - self.max_poc_lsb
64        } else {
65            self.prev_poc_msb
66        };
67
68        let poc = poc_msb + poc_lsb;
69
70        // Only reference pictures participate in the prevPicOrderCnt chain.
71        if nal_ref_idc != 0 {
72            self.prev_poc_msb = poc_msb;
73            self.prev_poc_lsb = poc_lsb;
74        }
75
76        poc
77    }
78}
79
80/// How each frame's picture order count is obtained, selected from the SPS's
81/// `pic_order_cnt_type`.
82pub(crate) enum PocStrategy {
83    /// `pic_order_cnt_type == 0`: read `pic_order_cnt_lsb` from every slice and
84    /// unwrap it (this is the type that can carry B-frame reordering).
85    FromSliceLsb(PocDecoder),
86    /// `pic_order_cnt_type == 2`: the bitstream guarantees decode order equals
87    /// display order (ITU-T H.264 §8.2.1.3 — no reordering is possible), so the
88    /// POC simply follows decode order.
89    DecodeOrder { next: i64 },
90}
91
92/// Determine the [`PocStrategy`] from an SPS's `pic_order_cnt_type`. Types 0 and
93/// 2 are supported; type 1 (rare, delta-based) is not.
94pub(crate) fn strategy_from_sps(sps: &SeqParameterSet) -> Result<PocStrategy> {
95    match sps.pic_order_cnt {
96        PicOrderCntType::TypeZero {
97            log2_max_pic_order_cnt_lsb_minus4,
98        } => Ok(PocStrategy::FromSliceLsb(PocDecoder::new(
99            log2_max_pic_order_cnt_lsb_minus4,
100        ))),
101        PicOrderCntType::TypeTwo => Ok(PocStrategy::DecodeOrder { next: 0 }),
102        PicOrderCntType::TypeOne { .. } => Err(Error::H264Poc(
103            "uses pic_order_cnt_type 1, which is not supported".to_string(),
104        )),
105    }
106}
107
108/// Parse an SPS NAL and determine its [`PocStrategy`].
109pub(crate) fn parse_sps(nal: &RefNal<'_>) -> Result<(SeqParameterSet, PocStrategy)> {
110    let sps = SeqParameterSet::from_bits(nal.rbsp_bits())
111        .map_err(|e| Error::H264Poc(format!("bad SPS: {e:?}")))?;
112    let strategy = strategy_from_sps(&sps)?;
113    Ok((sps, strategy))
114}
115
116/// Extract `(is_idr, nal_ref_idc, pic_order_cnt_lsb)` from the first slice NAL
117/// unit in a decoded sample.
118fn read_slice_poc_lsb(ctx: &H264ParsingContext, nals: &[Vec<u8>]) -> Result<(bool, u8, i64)> {
119    for nal_bytes in nals {
120        let nal = RefNal::new(nal_bytes, &[], true);
121        let header = nal
122            .header()
123            .map_err(|e| Error::H264Poc(format!("bad NAL header: {e:?}")))?;
124        let unit_type = header.nal_unit_type();
125        if !matches!(
126            unit_type,
127            UnitType::SliceLayerWithoutPartitioningIdr
128                | UnitType::SliceLayerWithoutPartitioningNonIdr
129        ) {
130            continue;
131        }
132        let is_idr = unit_type == UnitType::SliceLayerWithoutPartitioningIdr;
133        let mut r = nal.rbsp_bits();
134        let (slice_header, _sps, _pps) = SliceHeader::from_bits(ctx, &mut r, header)
135            .map_err(|e| Error::H264Poc(format!("bad slice header: {e:?}")))?;
136        let poc_lsb = match slice_header.pic_order_cnt_lsb {
137            Some(PicOrderCountLsb::Frame(lsb)) => lsb as i64,
138            Some(_) => {
139                return Err(Error::H264Poc(
140                    "field pictures are not supported".to_string(),
141                ));
142            }
143            None => {
144                return Err(Error::H264Poc(
145                    "slice has no pic_order_cnt_lsb (unsupported pic_order_cnt_type)".to_string(),
146                ));
147            }
148        };
149        return Ok((is_idr, header.nal_ref_idc(), poc_lsb));
150    }
151    Err(Error::H264Poc("sample has no slice NAL unit".to_string()))
152}
153
154/// Advance a [`PocStrategy`] by one frame (whose NAL units are `nals`, in decode
155/// order) and return that frame's POC. `ctx` must already contain the SPS/PPS
156/// referenced by the slice.
157pub(crate) fn advance_poc(
158    strategy: &mut PocStrategy,
159    ctx: &H264ParsingContext,
160    nals: &[Vec<u8>],
161) -> Result<i64> {
162    match strategy {
163        PocStrategy::FromSliceLsb(decoder) => {
164            let (is_idr, nal_ref_idc, poc_lsb) = read_slice_poc_lsb(ctx, nals)?;
165            Ok(decoder.next_poc(is_idr, nal_ref_idc, poc_lsb))
166        }
167        PocStrategy::DecodeOrder { next } => {
168            let poc = *next;
169            *next += 1;
170            Ok(poc)
171        }
172    }
173}
174
175/// Accumulates the H.264 parsing context (SPS/PPS) and the POC strategy as
176/// samples are read, so a whole file can be walked (in decode order) and each
177/// frame's picture order count reconstructed the same way.
178#[derive(Default)]
179pub struct PocReader {
180    ctx: H264ParsingContext,
181    // The POC strategy comes from the SPS's `pic_order_cnt_type`, so it can only
182    // be chosen once an SPS has been seen. MP4 keeps SPS/PPS in the container;
183    // Annex B streams carry them inline (picked up per-frame).
184    strategy: Option<PocStrategy>,
185}
186
187impl PocReader {
188    pub fn new() -> Self {
189        Self::default()
190    }
191
192    /// Record an SPS: feed it to the parsing context and, on the first one, fix
193    /// the POC strategy from its `pic_order_cnt_type`.
194    fn put_sps(&mut self, nal: &RefNal<'_>) -> Result<()> {
195        let (sps, strategy) = parse_sps(nal)?;
196        if self.strategy.is_none() {
197            self.strategy = Some(strategy);
198        }
199        self.ctx.put_seq_param_set(sps);
200        Ok(())
201    }
202
203    /// Seed SPS/PPS from container-level metadata (MP4). No-op for Annex B,
204    /// which carries them inline (handled in [`Self::poc_for_frame`]).
205    pub fn seed_from_container<H: SeekableH264Source>(
206        &mut self,
207        src: &H264Source<H>,
208    ) -> Result<()> {
209        if let Some(sps_bytes) = src.as_seekable_h264_source().first_sps() {
210            let nal = RefNal::new(&sps_bytes, &[], true);
211            self.put_sps(&nal)?;
212        }
213        if let Some(pps_bytes) = src.as_seekable_h264_source().first_pps() {
214            let nal = RefNal::new(&pps_bytes, &[], true);
215            let pps = PicParameterSet::from_bits(&self.ctx, nal.rbsp_bits())
216                .map_err(|e| Error::H264Poc(format!("bad PPS: {e:?}")))?;
217            self.ctx.put_pic_param_set(pps);
218        }
219        Ok(())
220    }
221
222    /// Feed any inline SPS/PPS (Annex B) carried with this frame, then return
223    /// the frame's POC.
224    pub fn poc_for_frame(&mut self, nals: &[Vec<u8>]) -> Result<i64> {
225        for nal_bytes in nals {
226            let nal = RefNal::new(nal_bytes, &[], true);
227            let Ok(header) = nal.header() else { continue };
228            match header.nal_unit_type() {
229                UnitType::SeqParameterSet => self.put_sps(&nal)?,
230                UnitType::PicParameterSet => {
231                    let pps = PicParameterSet::from_bits(&self.ctx, nal.rbsp_bits())
232                        .map_err(|e| Error::H264Poc(format!("bad PPS: {e:?}")))?;
233                    self.ctx.put_pic_param_set(pps);
234                }
235                _ => {}
236            }
237        }
238        // `self.ctx` and `self.strategy` are disjoint fields, so the immutable
239        // borrow of the context and the mutable borrow of the strategy coexist.
240        match &mut self.strategy {
241            None => Err(Error::H264Poc(
242                "slice data appeared before any SPS".to_string(),
243            )),
244            Some(strategy) => advance_poc(strategy, &self.ctx, nals),
245        }
246    }
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252
253    #[test]
254    fn parse_sps_selects_strategy_by_pic_order_cnt_type() {
255        // Real SPS NAL units (EBSP, including the NAL header byte) captured from
256        // sample recordings. `pic_order_cnt_type == 0` (explicit poc_lsb, can
257        // carry B-frame reordering) vs `== 2` (decode order == display order).
258        const SPS_TYPE0: &[u8] = &[
259            0x67, 0xf4, 0x00, 0x28, 0x91, 0x9b, 0x28, 0x0f, 0x00, 0x44, 0xfc, 0x4c, 0xd9, 0x00,
260            0x00, 0x03, 0x00, 0x01, 0x00, 0x00, 0x03, 0x00, 0x32, 0x0f, 0x18, 0x31, 0x96,
261        ];
262        const SPS_TYPE2: &[u8] = &[
263            0x67, 0x64, 0x44, 0x28, 0xac, 0x4d, 0x00, 0xf0, 0x04, 0x4f, 0xcb, 0x34, 0xb7, 0x00,
264            0x00, 0x03, 0x00, 0x01, 0x00, 0x00, 0x03, 0x00, 0x3c, 0x0f, 0x08, 0x84, 0x6a,
265        ];
266        let (_, s0) = parse_sps(&RefNal::new(SPS_TYPE0, &[], true)).unwrap();
267        assert!(
268            matches!(s0, PocStrategy::FromSliceLsb(_)),
269            "pic_order_cnt_type 0 should read poc_lsb from slices"
270        );
271        let (_, s2) = parse_sps(&RefNal::new(SPS_TYPE2, &[], true)).unwrap();
272        assert!(
273            matches!(s2, PocStrategy::DecodeOrder { .. }),
274            "pic_order_cnt_type 2 should fall back to decode order"
275        );
276    }
277
278    #[test]
279    fn poc_decoder_handles_simple_ipbb_gop() {
280        let mut dec = PocDecoder::new(4); // MaxPicOrderCntLsb = 256
281        // I(ref), P(ref), B(non-ref), B(non-ref), repeating POC pattern
282        // typical of an IBBP-style GOP with POC step 2 per displayed frame.
283        assert_eq!(dec.next_poc(true, 1, 0), 0); // I, poc 0
284        assert_eq!(dec.next_poc(false, 1, 6), 6); // P, poc 6
285        assert_eq!(dec.next_poc(false, 0, 2), 2); // B, poc 2
286        assert_eq!(dec.next_poc(false, 0, 4), 4); // B, poc 4
287    }
288
289    #[test]
290    fn poc_decoder_unwraps_lsb_wraparound() {
291        let mut dec = PocDecoder::new(0); // MaxPicOrderCntLsb = 16
292        // Step by 2 each reference frame, staying well under
293        // MaxPicOrderCntLsb/2 (8) per step so no wrap is triggered yet.
294        assert_eq!(dec.next_poc(true, 1, 0), 0);
295        assert_eq!(dec.next_poc(false, 1, 2), 2);
296        assert_eq!(dec.next_poc(false, 1, 4), 4);
297        assert_eq!(dec.next_poc(false, 1, 6), 6);
298        assert_eq!(dec.next_poc(false, 1, 8), 8);
299        assert_eq!(dec.next_poc(false, 1, 10), 10);
300        assert_eq!(dec.next_poc(false, 1, 12), 12);
301        assert_eq!(dec.next_poc(false, 1, 14), 14);
302        // lsb wraps from 14 back down to 0. The raw backward delta (14)
303        // meets MaxPicOrderCntLsb/2, so this is really a forward step to
304        // poc 16, not a jump back near zero.
305        assert_eq!(dec.next_poc(false, 1, 0), 16);
306    }
307}
frame_source/h264_poc.rs

frame_source/
h264_poc.rs
