h264: fix 'invalid POC' error

related to

bluenviron/mediamtx#1552
bluenviron/mediamtx#1002

pkg/codecs/h264/dts_extractor.go

@@ -92,7 +92,7 @@ func NewDTSExtractor() *DTSExtractor {
 	}
 }
 
-func (d *DTSExtractor) extractInner(au [][]byte, pts time.Duration) (time.Duration, error) {
+func (d *DTSExtractor) extractInner(au [][]byte, pts time.Duration) (time.Duration, bool, error) {
 	idrPresent := false
 
 	for _, nalu := range au {
@@ -103,7 +103,7 @@ func (d *DTSExtractor) extractInner(au [][]byte, pts time.Duration) (time.Durati
 				var spsp SPS
 				err := spsp.Unmarshal(nalu)
 				if err != nil {
-					return 0, fmt.Errorf("invalid SPS: %v", err)
+					return 0, false, fmt.Errorf("invalid SPS: %v", err)
 				}
 				d.sps = nalu
 				d.spsp = &spsp
@@ -119,40 +119,39 @@ func (d *DTSExtractor) extractInner(au [][]byte, pts time.Duration) (time.Durati
 	}
 
 	if d.spsp == nil {
-		return 0, fmt.Errorf("SPS not received yet")
+		return 0, false, fmt.Errorf("SPS not received yet")
 	}
 
 	if d.spsp.PicOrderCntType == 2 {
-		return pts, nil
+		return pts, false, nil
 	}
 
 	if d.spsp.PicOrderCntType == 1 {
-		return 0, fmt.Errorf("pic_order_cnt_type = 1 is not supported yet")
+		return 0, false, fmt.Errorf("pic_order_cnt_type = 1 is not supported yet")
 	}
 
 	if idrPresent {
 		d.expectedPOC = 0
-	} else {
-		d.expectedPOC += uint32(d.pocIncrement)
-		d.expectedPOC &= ((1 << (d.spsp.Log2MaxPicOrderCntLsbMinus4 + 4)) - 1)
-	}
-
-	if d.pauseDTS > 0 {
-		d.pauseDTS--
-		return d.prevDTS + 1*time.Millisecond, nil
-	}
+		d.pauseDTS = 0
 
-	if idrPresent {
 		if !d.prevDTSFilled || d.reorderedFrames == 0 {
-			return pts, nil
+			return pts, false, nil
 		}
 
-		return d.prevDTS + (pts-d.prevDTS)/time.Duration(d.reorderedFrames+1), nil
+		return d.prevDTS + (pts-d.prevDTS)/time.Duration(d.reorderedFrames+1), false, nil
+	}
+
+	d.expectedPOC += uint32(d.pocIncrement)
+	d.expectedPOC &= ((1 << (d.spsp.Log2MaxPicOrderCntLsbMinus4 + 4)) - 1)
+
+	if d.pauseDTS > 0 {
+		d.pauseDTS--
+		return d.prevDTS + 1*time.Millisecond, true, nil
 	}
 
 	poc, err := findPictureOrderCount(au, d.spsp)
 	if err != nil {
-		return 0, err
+		return 0, false, err
 	}
 
 	if d.pocIncrement == 2 && (poc%2) != 0 {
@@ -163,32 +162,46 @@ func (d *DTSExtractor) extractInner(au [][]byte, pts time.Duration) (time.Durati
 	pocDiff := int(getPictureOrderCountDiff(poc, d.expectedPOC, d.spsp)) + d.reorderedFrames*d.pocIncrement
 
 	if pocDiff < 0 {
-		return 0, fmt.Errorf("invalid POC")
+		if pocDiff < -20 {
+			return 0, false, fmt.Errorf("POC difference between frames is too big (%d)", pocDiff)
+		}
+
+		// this happens when there are B-frames immediately following an IDR frame
+		d.reorderedFrames -= pocDiff
+		d.pauseDTS = -pocDiff
+		return d.prevDTS + 1*time.Millisecond, true, nil
 	}
 
 	if pocDiff == 0 {
-		return pts, nil
+		return pts, false, nil
+	}
+
+	if pocDiff > 20 {
+		return 0, false, fmt.Errorf("POC difference between frames is too big (%d)", pocDiff)
 	}
 
-	reorderedFrames := (pocDiff - d.reorderedFrames*d.pocIncrement) / d.pocIncrement
+	reorderedFrames := (pocDiff)/d.pocIncrement - d.reorderedFrames
 	if reorderedFrames > d.reorderedFrames {
+		// reordered frames detected, add them to the count and pause DTS
 		d.pauseDTS = (reorderedFrames - d.reorderedFrames - 1)
 		d.reorderedFrames = reorderedFrames
-		return d.prevDTS + 1*time.Millisecond, nil
+		return d.prevDTS + 1*time.Millisecond, false, nil
 	}
 
-	return d.prevDTS + (pts-d.prevDTS)*time.Duration(d.pocIncrement)/time.Duration(pocDiff+d.pocIncrement), nil
+	return d.prevDTS + (pts-d.prevDTS)*time.Duration(d.pocIncrement)/time.Duration(pocDiff+d.pocIncrement), false, nil
 }
 
-// Extract extracts the DTS of a access unit.
+// Extract extracts the DTS of an access unit.
 func (d *DTSExtractor) Extract(au [][]byte, pts time.Duration) (time.Duration, error) {
-	dts, err := d.extractInner(au, pts)
+	dts, skipChecks, err := d.extractInner(au, pts)
 	if err != nil {
 		return 0, err
 	}
 
-	if dts > pts {
-		return 0, fmt.Errorf("DTS is greater than PTS")
+	if !skipChecks {
+		if dts > pts {
+			return 0, fmt.Errorf("DTS is greater than PTS")
+		}
 	}
 
 	if d.prevDTSFilled && dts <= d.prevDTS {

pkg/codecs/h264/dts_extractor_test.go

@@ -8,19 +8,19 @@ import (
 )
 
 func TestDTSExtractor(t *testing.T) {
-	type sequenceSample struct {
+	type sample struct {
 		au  [][]byte
 		dts time.Duration
 		pts time.Duration
 	}
 
 	for _, ca := range []struct {
 		name     string
-		sequence []sequenceSample
+		sequence []sample
 	}{
 		{
 			"with timing info",
-			[]sequenceSample{
+			[]sample{
 				{
 					[][]byte{
 						{ // SPS
@@ -84,7 +84,7 @@ func TestDTSExtractor(t *testing.T) {
 		},
 		{
 			"no timing info",
-			[]sequenceSample{
+			[]sample{
 				{
 					[][]byte{
 						{ // SPS
@@ -164,7 +164,7 @@ func TestDTSExtractor(t *testing.T) {
 		},
 		{
 			"poc increment = 1",
-			[]sequenceSample{
+			[]sample{
 				{
 					[][]byte{
 						{ // SPS
@@ -197,6 +197,48 @@ func TestDTSExtractor(t *testing.T) {
 				},
 			},
 		},
+		{
+			"B-frames after IDR (OBS 29.1.3 QuickSync on Windows)",
+			[]sample{
+				{
+					[][]byte{
+						{ // SPS
+							0x27, 0x64, 0x00, 0x2a, 0xac, 0x2d, 0x90, 0x07,
+							0x80, 0x22, 0x7e, 0x5c, 0x05, 0xa8, 0x08, 0x08,
+							0x0a, 0x00, 0x00, 0x03, 0x00, 0x02, 0x00, 0x00,
+							0x03, 0x00, 0xf1, 0xd0, 0x80, 0x04, 0xc4, 0x80,
+							0x00, 0x09, 0x89, 0x68, 0xde, 0xf7, 0xc1, 0xda,
+							0x1c, 0x31, 0x92,
+						},
+						{ // IDR
+							0x65, 0x88, 0x80, 0x14, 0x3, 0xff, 0xde, 0x8, 0xe4, 0x74,
+						},
+					},
+					1916 * time.Millisecond,
+					1916 * time.Millisecond,
+				},
+				{ // b-frame
+					[][]byte{{0x41, 0x9e, 0x3, 0xe4, 0x3f, 0x0, 0x0, 0x3, 0x0, 0x0}},
+					1917 * time.Millisecond,
+					1883 * time.Millisecond,
+				},
+				{ // b-frame
+					[][]byte{{0x1, 0x9e, 0x5, 0xd4, 0x7f, 0x0, 0x0, 0x3, 0x0, 0x0}},
+					1918 * time.Millisecond,
+					1867 * time.Millisecond,
+				},
+				{ // p-frame
+					[][]byte{{0x1, 0x9e, 0x5, 0xf4, 0x7f, 0x0, 0x0, 0x3, 0x0, 0x0}},
+					1919 * time.Millisecond,
+					1899 * time.Millisecond,
+				},
+				{ // p-frame
+					[][]byte{{0x1, 0x9e, 0x5, 0xf4, 0x7f, 0x0, 0x0, 0x3, 0x0, 0x0}},
+					1920 * time.Millisecond,
+					1983 * time.Millisecond,
+				},
+			},
+		},
 	} {
 		t.Run(ca.name, func(t *testing.T) {
 			ex := NewDTSExtractor()