h264: fix DTS extraction in case of long NALUs (#98)

* add edgecase for DTSExtractor * apply review * improve performance --------- Co-authored-by: aler9 <46489434+aler9@users.noreply.github.com>
bluenviron · Nov 17, 2023 · 151a3cb · 151a3cb
1 parent 1b36c5c
commit 151a3cb
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 5 deletions.
diff --git a/pkg/codecs/h264/dts_extractor.go b/pkg/codecs/h264/dts_extractor.go
@@ -10,16 +10,24 @@ import (
 
 const (
  maxReorderedFrames = 10
+ /*
+ (max_size(first_mb_in_slice) + max_size(slice_type) + max_size(pic_parameter_set_id) +
+ max_size(frame_num) + max_size(pic_order_cnt_lsb)) * 4 / 3 =
+ (3 * max_size(golomb) + (max(Log2MaxFrameNumMinus4) + 4) / 8 + (max(Log2MaxPicOrderCntLsbMinus4) + 4) / 8) * 4 / 3 =
+ (3 * 4 + 2 + 2) * 4 / 3 = 22
+ */
+ maxBytesToGetPOC = 22
 )
 
 func getPictureOrderCount(buf []byte, sps *SPS) (uint32, error) {
- if len(buf) < 6 {
- return 0, fmt.Errorf("not enough bits")
- }
+ buf = buf[1:]
+ lb := len(buf)
 
- buf = EmulationPreventionRemove(buf[:6])
+ if lb > maxBytesToGetPOC {
+ lb = maxBytesToGetPOC
+ }
 
- buf = buf[1:]
+ buf = EmulationPreventionRemove(buf[:lb])
  pos := 0
 
  _, err := bits.ReadGolombUnsigned(buf, &pos) // first_mb_in_slice

diff --git a/pkg/codecs/h264/dts_extractor_test.go b/pkg/codecs/h264/dts_extractor_test.go
@@ -278,6 +278,33 @@ func TestDTSExtractor(t *testing.T) {
  },
  },
  },
+ {
+ "Log2MaxPicOrderCntLsbMinus4 = 12",
+ []sample{
+ {
+ [][]byte{
+ { // SPS
+ 0x67, 0x42, 0xc0, 0x1e, 0x8c, 0x8d, 0x40, 0x50, 0x17, 0xfc, 0xb0, 0x0f, 0x08, 0x84, 0x6a,
+ },
+ { // PPS
+ 0x68, 0xce, 0x3c, 0x80,
+ },
+ { // IDR
+ 0x65, 0x88, 0x80, 0x14, 0x3, 0xff, 0xde, 0x8, 0xe4, 0x74,
+ },
+ },
+ 0,
+ 0,
+ },
+ {
+ [][]byte{
+ {0x61, 0x00, 0xf0, 0xe0, 0x00, 0x40, 0x00, 0xbe, 0x47, 0x9b}, // non-IDR
+ },
+ 40 * time.Millisecond,
+ 40 * time.Millisecond,
+ },
+ },
+ },
  } {
  t.Run(ca.name, func(t *testing.T) {
  ex := NewDTSExtractor()