h264: support extracting DTS from nvenc (https://github.com/aler9/rtsp-simple-server/issues/989)

pkg/h264/dtsextractor.go

@@ -68,7 +68,7 @@ func getPOC(buf []byte, sps *SPS) (uint32, error) {
 	return uint32(picOrderCntLsb), nil
 }
 
-func getNALUSPOC(nalus [][]byte, sps *SPS) (uint32, error) {
+func findPOC(nalus [][]byte, sps *SPS) (uint32, error) {
 	for _, nalu := range nalus {
 		typ := NALUType(nalu[0] & 0x1F)
 		if typ == NALUTypeIDR || typ == NALUTypeNonIDR {
@@ -94,6 +94,70 @@ func getPOCDiff(poc1 uint32, poc2 uint32, sps *SPS) int32 {
 	return diff
 }
 
+func getSEIPicTimingDPBOutputDelay(buf []byte, sps *SPS) (uint32, bool) {
+	buf = AntiCompetitionRemove(buf)
+	pos := 1
+
+	for {
+		if pos >= (len(buf) - 1) {
+			return 0, false
+		}
+
+		payloadType := 0
+		for {
+			byt := buf[pos]
+			pos++
+			payloadType += int(byt)
+			if byt != 0xFF {
+				break
+			}
+		}
+
+		payloadSize := 0
+		for {
+			byt := buf[pos]
+			pos++
+			payloadSize += int(byt)
+			if byt != 0xFF {
+				break
+			}
+		}
+
+		if payloadType == 1 { // timing info
+			br := bitio.NewReader(bytes.NewReader(buf[pos : pos+payloadSize]))
+
+			// cpbRemovalDelay
+			_, err := br.ReadBits(sps.VUI.NalHRD.CpbRemovalDelayLengthMinus1 + 1)
+			if err != nil {
+				return 0, false
+			}
+
+			tmp, err := br.ReadBits(sps.VUI.NalHRD.DpbOutputDelayLengthMinus1 + 1)
+			if err != nil {
+				return 0, false
+			}
+			dpbOutputDelay := uint32(tmp)
+
+			return dpbOutputDelay, true
+		}
+
+		pos += payloadSize
+	}
+}
+
+func findSEIPicTimingDPBOutputDelay(nalus [][]byte, sps *SPS) (uint32, bool) {
+	for _, nalu := range nalus {
+		typ := NALUType(nalu[0] & 0x1F)
+		if typ == NALUTypeSEI {
+			ret, ok := getSEIPicTimingDPBOutputDelay(nalu, sps)
+			if ok {
+				return ret, true
+			}
+		}
+	}
+	return 0, false
+}
+
 // DTSExtractor is a utility that allows to extract NALU DTS from PTS.
 type DTSExtractor struct {
 	sps          []byte
@@ -110,27 +174,22 @@ func NewDTSExtractor() *DTSExtractor {
 	return &DTSExtractor{}
 }
 
-func (d *DTSExtractor) extractInner(
+func (d *DTSExtractor) extractInner(nalus [][]byte, pts time.Duration) (time.Duration, int32, error) {
-	nalus [][]byte,
-	pts time.Duration,
-) (time.Duration, int32, error) {
 	idrPresent := false
 
 	for _, nalu := range nalus {
 		typ := NALUType(nalu[0] & 0x1F)
 		switch typ {
-		// parse SPS
 		case NALUTypeSPS:
 			if d.sps == nil || !bytes.Equal(d.sps, nalu) {
 				var spsp SPS
 				err := spsp.Unmarshal(nalu)
 				if err != nil {
-					return 0, 0, err
+					return 0, 0, fmt.Errorf("invalid SPS: %v", err)
 				}
 				d.sps = append([]byte(nil), nalu...)
 				d.spsp = &spsp
 
-				// in case of B-frames, we have to subtract from DTS the maximum number of reordered frames
 				if d.spsp.VUI != nil && d.spsp.VUI.TimingInfo != nil &&
 					d.spsp.VUI.BitstreamRestriction != nil {
 					d.ptsDTSOffset = time.Duration(math.Round(float64(
@@ -141,7 +200,6 @@ func (d *DTSExtractor) extractInner(
 				}
 			}
 
-		// set IDR present flag
 		case NALUTypeIDR:
 			idrPresent = true
 		}
@@ -151,7 +209,24 @@ func (d *DTSExtractor) extractInner(
 		return 0, 0, fmt.Errorf("SPS not received yet")
 	}
 
-	if idrPresent || d.spsp.PicOrderCntType == 2 {
+	switch {
+	// DTS is computed from SEI
+	case d.spsp.VUI != nil && d.spsp.VUI.TimingInfo != nil && d.spsp.VUI.NalHRD != nil:
+		dpbOutputDelay, ok := findSEIPicTimingDPBOutputDelay(nalus, d.spsp)
+		if !ok {
+			return 0, 0, fmt.Errorf("SEI / pic_timing not found")
+		}
+
+		return pts - time.Duration(dpbOutputDelay)/2*time.Second*
+			time.Duration(d.spsp.VUI.TimingInfo.NumUnitsInTick)*2/time.Duration(d.spsp.VUI.TimingInfo.TimeScale), 0, nil
+
+	// DTS is always equal to PTS
+	case d.spsp.PicOrderCntType == 2:
+		return pts, 0, nil
+
+	// DTS is computed by using POC, timing infos and max_num_reorder_frames
+	case d.spsp.VUI != nil && d.spsp.VUI.TimingInfo != nil && d.spsp.VUI.BitstreamRestriction != nil:
+		if idrPresent {
 			d.expectedPOC = 0
 			return pts - d.ptsDTSOffset, 0, nil
 		}
@@ -160,7 +235,7 @@ func (d *DTSExtractor) extractInner(
 		d.expectedPOC += 2
 		d.expectedPOC &= ((1 << (d.spsp.Log2MaxPicOrderCntLsbMinus4 + 4)) - 1)
 
-	poc, err := getNALUSPOC(nalus, d.spsp)
+		poc, err := findPOC(nalus, d.spsp)
 		if err != nil {
 			return 0, 0, err
 		}
@@ -172,8 +247,7 @@ func (d *DTSExtractor) extractInner(
 		}
 
 		// special case to eliminate errors near 0
-	if d.spsp.VUI != nil && d.spsp.VUI.TimingInfo != nil && d.spsp.VUI.BitstreamRestriction != nil &&
+		if pocDiff == -int32(d.spsp.VUI.BitstreamRestriction.MaxNumReorderFrames)*2 {
-		pocDiff == -int32(d.spsp.VUI.BitstreamRestriction.MaxNumReorderFrames)*2 {
 			return pts, pocDiff, nil
 		}
 
@@ -189,13 +263,14 @@ func (d *DTSExtractor) extractInner(
 		// pocDiff : prevPOCDiff = (pts - dts - ptsDTSOffset) : (prevPTS - prevDTS - ptsDTSOffset)
 		return pts - d.ptsDTSOffset + time.Duration(math.Round(float64(d.prevDTS-d.prevPTS+d.ptsDTSOffset)*
 			float64(pocDiff)/float64(d.prevPOCDiff))), pocDiff, nil
+
+	default:
+		return 0, 0, fmt.Errorf("unable to compute H264 DTS")
+	}
 }
 
-// Extract extracts the DTS of a NALU group.
+// Extract extracts the DTS of a group of NALUs.
-func (d *DTSExtractor) Extract(
+func (d *DTSExtractor) Extract(nalus [][]byte, pts time.Duration) (time.Duration, error) {
-	nalus [][]byte,
-	pts time.Duration,
-) (time.Duration, error) {
 	dts, pocDiff, err := d.extractInner(nalus, pts)
 	if err != nil {
 		return 0, err

pkg/h264/dtsextractor_test.go

@@ -8,11 +8,19 @@ import (
 )
 
 func TestDTSExtractor(t *testing.T) {
-	sequence := []struct {
+	type sequenceSample struct {
 		nalus [][]byte
 		pts   time.Duration
 		dts   time.Duration
+	}
+
+	for _, ca := range []struct {
+		name     string
+		sequence []sequenceSample
 	}{
+		{
+			"max_num_reorder_frames-based",
+			[]sequenceSample{
 				{
 					[][]byte{
 						{
@@ -132,13 +140,63 @@ func TestDTSExtractor(t *testing.T) {
 					200 * time.Millisecond,
 					200 * time.Millisecond,
 				},
-	}
+			},
-
+		},
+		{
+			"sei-based",
+			[]sequenceSample{
+				{
+					[][]byte{
+						// SEI (buffering period)
+						{6, 0, 7, 128, 117, 48, 0, 0, 3, 0, 64, 128},
+						// SEI (pic timing)
+						{6, 1, 4, 0, 0, 8, 16, 128},
+						// SPS
+						{
+							103, 100, 0, 42, 172, 44, 172, 7,
+							128, 34, 126, 92, 5, 168, 8, 8,
+							10, 0, 0, 7, 208, 0, 3, 169,
+							129, 192, 0, 0, 76, 75, 0, 0,
+							38, 37, 173, 222, 92, 20,
+						},
+					},
+					0,
+					-16666666 * time.Nanosecond,
+				},
+				{
+					[][]byte{
+						// SEI
+						{6, 1, 4, 0, 2, 32, 16, 128},
+					},
+					66666666 * time.Nanosecond,
+					0,
+				},
+				{
+					[][]byte{
+						// SEI
+						{6, 1, 4, 0, 4, 0, 16, 128},
+					},
+					16666666 * time.Nanosecond,
+					16666666 * time.Nanosecond,
+				},
+				{
+					[][]byte{
+						// SEI
+						{6, 1, 4, 0, 6, 0, 16, 128},
+					},
+					33333333 * time.Nanosecond,
+					33333333 * time.Nanosecond,
+				},
+			},
+		},
+	} {
+		t.Run(ca.name, func(t *testing.T) {
 			ex := NewDTSExtractor()
-
+			for _, sample := range ca.sequence {
-	for _, sample := range sequence {
 				dts, err := ex.Extract(sample.nalus, sample.pts)
 				require.NoError(t, err)
 				require.Equal(t, sample.dts, dts)
 			}
+		})
+	}
 }

pkg/h264/sps.go

@@ -841,7 +841,7 @@ func (s SPS) Height() int {
 	return int((2 - f) * (s.PicHeightInMbsMinus1 + 1) * 16)
 }
 
-// FPS returns the frame per second of the video.
+// FPS returns the frames per second of the video.
 func (s SPS) FPS() float64 {
 	if s.VUI == nil || s.VUI.TimingInfo == nil {
 		return 0