h264: fix DTS extraction of NVC4000 streams (#136)

(bluenviron/mediamtx#3094)
bluenviron · Jul 1, 2024 · bdc22be · bdc22be
1 parent 2fe38f4
commit bdc22be
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 6 deletions.
diff --git a/pkg/codecs/h264/dts_extractor.go b/pkg/codecs/h264/dts_extractor.go
@@ -19,7 +19,7 @@ const (
 	maxBytesToGetPOC = 22
 )
 
-func getPictureOrderCount(buf []byte, sps *SPS) (uint32, error) {
+func getPictureOrderCount(buf []byte, sps *SPS, idr bool) (uint32, error) {
 	buf = buf[1:]
 	lb := len(buf)
 
@@ -50,6 +50,13 @@ func getPictureOrderCount(buf []byte, sps *SPS) (uint32, error) {
 		return 0, err
 	}
 
+	if idr {
+		_, err = bits.ReadGolombUnsigned(buf, &pos) // idr_pic_id
+		if err != nil {
+			return 0, err
+		}
+	}
+
 	picOrderCntLsb, err := bits.ReadBits(buf, &pos, int(sps.Log2MaxPicOrderCntLsbMinus4+4))
 	if err != nil {
 		return 0, err
@@ -130,9 +137,14 @@ func (d *DTSExtractor) extractInner(au [][]byte, pts time.Duration) (time.Durati
 
 	switch {
 	case idr != nil:
-		d.expectedPOC = 0
 		d.pauseDTS = 0
 
+		var err error
+		d.expectedPOC, err = getPictureOrderCount(idr, d.spsp, true)
+		if err != nil {
+			return 0, false, err
+		}
+
 		if !d.prevDTSFilled || d.reorderedFrames == 0 {
 			return pts, false, nil
 		}
@@ -148,7 +160,7 @@ func (d *DTSExtractor) extractInner(au [][]byte, pts time.Duration) (time.Durati
 			return d.prevDTS + 1*time.Millisecond, true, nil
 		}
 
-		poc, err := getPictureOrderCount(nonIDR, d.spsp)
+		poc, err := getPictureOrderCount(nonIDR, d.spsp, false)
 		if err != nil {
 			return 0, false, err
 		}

diff --git a/pkg/codecs/h264/dts_extractor_test.go b/pkg/codecs/h264/dts_extractor_test.go
@@ -173,7 +173,7 @@ var casesDTSExtractor = []struct {
 						0x02,
 					},
 					{ // IDR
-						0x65, 0xb8, 0x00, 0x00, 0x0b, 0xc8,
+						0x65, 0xb8, 0x00, 0x00, 0x0b, 0xc8, 0x00, 0x00, 0x00,
 					},
 				},
 				61 * time.Millisecond,
@@ -289,7 +289,7 @@ var casesDTSExtractor = []struct {
 						0x68, 0xce, 0x3c, 0x80,
 					},
 					{ // IDR
-						0x65, 0x88, 0x80, 0x14, 0x3, 0xff, 0xde, 0x8, 0xe4, 0x74,
+						0x65, 0x88, 0x80, 0x14, 0x3, 0x00, 0x00, 0x00, 0x00, 0x00,
 					},
 				},
 				0,
@@ -304,6 +304,73 @@ var casesDTSExtractor = []struct {
 			},
 		},
 	},
+	{
+		"issue mediamtx/3094 (non-zero IDR POC)",
+		[]sample{
+			{
+				[][]byte{
+					{
+						0x67, 0x42, 0x80, 0x28, 0x8c, 0x8d, 0x40, 0x5a,
+						0x09, 0x22,
+					},
+					{
+						0x68, 0xce, 0x3c, 0x80,
+					},
+					{
+						0x65, 0xb8, 0x00, 0x0c, 0xa2, 0x40, 0x33, 0x93,
+						0x14, 0x00, 0x04, 0x1a, 0x6d, 0x6d, 0x6d, 0x6d,
+						0x6d, 0x6d, 0x5d, 0xaa, 0xb5, 0xb5, 0xb5, 0xb5,
+						0xb5, 0xb5, 0xb5, 0xb5, 0xb5, 0xb5, 0xb5, 0xb5,
+						0xb5, 0xb5, 0x76, 0xb6, 0xb6, 0xb6, 0xaa, 0xd6,
+						0xd6, 0xd6, 0xd6, 0xd6, 0xd6, 0xd6, 0xd6, 0xd6,
+						0xd6, 0xd6, 0xd6, 0xd6, 0xd6, 0xd6, 0xd6, 0xd5,
+						0xda, 0xda, 0xaa, 0x7a, 0x7a, 0x7a, 0x7a, 0x7a,
+						0x79, 0x1e, 0xde, 0xde, 0xde, 0xde, 0xde, 0xde,
+						0xde, 0xde, 0xde, 0xde, 0xde, 0xde, 0xde, 0xde,
+						0xde, 0xde, 0xde, 0xde, 0xde, 0xde, 0xde, 0xde,
+						0xde, 0xde, 0xde, 0xde, 0xde, 0xde, 0xde, 0xde,
+						0xde, 0xde, 0xde, 0xde,
+					},
+				},
+				0,
+				0,
+			},
+			{
+				[][]byte{
+					{
+						0x41, 0xe0, 0x00, 0x65, 0x12, 0x80, 0xce, 0x78,
+						0x16, 0x00, 0x99, 0xff, 0xff, 0xff, 0xe0, 0xe4,
+						0x1a, 0x7f, 0xff, 0xff, 0xea, 0x11, 0x01, 0x01,
+						0xff, 0xff, 0xfc, 0x20, 0x08, 0x3f, 0xff, 0xff,
+						0xfc, 0x0f, 0x22, 0x7f, 0xff, 0xff, 0xff, 0xff,
+						0xff, 0xff, 0xff, 0xff, 0xb8, 0x60, 0x04, 0x87,
+						0x02, 0xc8, 0x18, 0x38, 0x60, 0x04, 0x87, 0x03,
+						0x00, 0x35, 0xa8, 0x16, 0x40, 0x9b, 0x04, 0xd0,
+						0x11, 0x00, 0x24, 0x38, 0x11, 0x01, 0x6c, 0x16,
+						0x41, 0x60, 0x2c, 0x82, 0xd8, 0x2c, 0x05, 0x90,
+						0x5b, 0x05, 0x80, 0xb2, 0x0b, 0x60, 0xb0, 0x16,
+						0x41, 0x6c, 0x16, 0x02, 0xc8, 0x2d, 0x82, 0xc0,
+						0x59, 0x05, 0xb0, 0x58,
+					},
+				},
+				120 * time.Millisecond,
+				120 * time.Millisecond,
+			},
+			{
+				[][]byte{
+					{
+						0x41, 0xe0, 0x00, 0xa5, 0x13, 0x00, 0xce, 0xf0,
+						0x2c, 0x70, 0x20, 0x01, 0x43, 0xc0, 0x8b, 0xc3,
+						0x01, 0x99, 0x60, 0x80, 0x04, 0x07, 0x06, 0x39,
+						0xe0, 0x80, 0x04, 0x04, 0x37, 0x80, 0x90, 0xe4,
+						0x06, 0x9c, 0xa0, 0x23, 0x60, 0x06, 0x25, 0x80,
+					},
+				},
+				160 * time.Millisecond,
+				160 * time.Millisecond,
+			},
+		},
+	},
 }
 
 func TestDTSExtractor(t *testing.T) {
@@ -350,7 +417,7 @@ func FuzzDTSExtractorSecondAU(f *testing.F) {
 				0xf8, 0x08,
 			},
 			{ // IDR
-				0x05,
+				0x25, 0xb8, 0x08, 0x02, 0x1f, 0xff,
 			},
 		}, 0)
 		require.NoError(t, err)