move DTS extractor into gortsplib

go.mod

@@ -5,13 +5,12 @@ go 1.17
 require (
 	code.cloudfoundry.org/bytefmt v0.0.0-20211005130812-5bb3c17173e5
 	github.com/abema/go-mp4 v0.7.2
-	github.com/aler9/gortsplib v0.0.0-20220529122539-7c2e3c03d1f4
+	github.com/aler9/gortsplib v0.0.0-20220602092441-d19c37025dc4
 	github.com/asticode/go-astits v1.10.1-0.20220319093903-4abe66a9b757
 	github.com/fsnotify/fsnotify v1.4.9
 	github.com/gin-gonic/gin v1.7.2
 	github.com/gookit/color v1.4.2
 	github.com/grafov/m3u8 v0.11.1
-	github.com/icza/bitio v1.1.0
 	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51
 	github.com/notedit/rtmp v0.0.0
 	github.com/orcaman/writerseeker v0.0.0
@@ -33,6 +32,7 @@ require (
 	github.com/go-playground/validator/v10 v10.4.1 // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/google/uuid v1.1.2 // indirect
+	github.com/icza/bitio v1.1.0 // indirect
 	github.com/json-iterator/go v1.1.9 // indirect
 	github.com/leodido/go-urn v1.2.0 // indirect
 	github.com/mattn/go-isatty v0.0.12 // indirect

go.sum

@@ -6,8 +6,8 @@ github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 h1:JYp7IbQjafo
 github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d h1:UQZhZ2O0vMHr2cI+DC1Mbh0TJxzA3RcLoMsFw+aXw7E=
 github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho=
-github.com/aler9/gortsplib v0.0.0-20220529122539-7c2e3c03d1f4 h1:5gd/tqApfqpfjCAlV4pxlwzWmRYMO+m4+RYGgSnSN3Q=
-github.com/aler9/gortsplib v0.0.0-20220529122539-7c2e3c03d1f4/go.mod h1:i1e4CEs42IrbidMUNTSNOKmeGPCOHVX9P3BvPxzyMtI=
+github.com/aler9/gortsplib v0.0.0-20220602092441-d19c37025dc4 h1:0xhKBSxHasMGLUgdE/kTw5QmywTcil4oQwh0+MmK+6o=
+github.com/aler9/gortsplib v0.0.0-20220602092441-d19c37025dc4/go.mod h1:i1e4CEs42IrbidMUNTSNOKmeGPCOHVX9P3BvPxzyMtI=
 github.com/aler9/rtmp v0.0.0-20210403095203-3be4a5535927 h1:95mXJ5fUCYpBRdSOnLAQAdJHHKxxxJrVCiaqDi965YQ=
 github.com/aler9/rtmp v0.0.0-20210403095203-3be4a5535927/go.mod h1:vzuE21rowz+lT1NGsWbreIvYulgBpCGnQyeTyFblUHc=
 github.com/aler9/writerseeker v0.0.0-20220601075008-6f0e685b9c82 h1:9WgSzBLo3a9ToSVV7sRTBYZ1GGOZUpq4+5H3SN0UZq4=

internal/hls/muxer_variant_fmp4.go

@@ -1,181 +1,22 @@
 package hls
 
 import (
-	"bytes"
-	"fmt"
-	"math"
 	"time"
 
 	"github.com/aler9/gortsplib"
-	"github.com/aler9/gortsplib/pkg/h264"
-	"github.com/icza/bitio"
 )
 
 const (
 	fmp4VideoTimescale = 90000
 )
 
-func readGolombUnsigned(br *bitio.Reader) (uint32, error) {
-	leadingZeroBits := uint32(0)
-
-	for {
-		b, err := br.ReadBits(1)
-		if err != nil {
-			return 0, err
-		}
-
-		if b != 0 {
-			break
-		}
-
-		leadingZeroBits++
-	}
-
-	codeNum := uint32(0)
-
-	for n := leadingZeroBits; n > 0; n-- {
-		b, err := br.ReadBits(1)
-		if err != nil {
-			return 0, err
-		}
-
-		codeNum |= uint32(b) << (n - 1)
-	}
-
-	codeNum = (1 << leadingZeroBits) - 1 + codeNum
-
-	return codeNum, nil
-}
-
-func getPOC(buf []byte, sps *h264.SPS) (uint32, error) {
-	buf = h264.AntiCompetitionRemove(buf[:10])
-
-	isIDR := h264.NALUType(buf[0]&0x1F) == h264.NALUTypeIDR
-
-	r := bytes.NewReader(buf[1:])
-	br := bitio.NewReader(r)
-
-	// first_mb_in_slice
-	_, err := readGolombUnsigned(br)
-	if err != nil {
-		return 0, err
-	}
-
-	// slice_type
-	_, err = readGolombUnsigned(br)
-	if err != nil {
-		return 0, err
-	}
-
-	// pic_parameter_set_id
-	_, err = readGolombUnsigned(br)
-	if err != nil {
-		return 0, err
-	}
-
-	// frame_num
-	_, err = br.ReadBits(uint8(sps.Log2MaxFrameNumMinus4 + 4))
-	if err != nil {
-		return 0, err
-	}
-
-	if !sps.FrameMbsOnlyFlag {
-		return 0, fmt.Errorf("unsupported")
-	}
-
-	if isIDR {
-		// idr_pic_id
-		_, err := readGolombUnsigned(br)
-		if err != nil {
-			return 0, err
-		}
-	}
-
-	var picOrderCntLsb uint64
-	switch {
-	case sps.PicOrderCntType == 0:
-		picOrderCntLsb, err = br.ReadBits(uint8(sps.Log2MaxPicOrderCntLsbMinus4 + 4))
-		if err != nil {
-			return 0, err
-		}
-
-	default:
-		return 0, fmt.Errorf("pic_order_cnt_type = 1 is unsupported")
-	}
-
-	return uint32(picOrderCntLsb), nil
-}
-
-func getNALUSPOC(nalus [][]byte, sps *h264.SPS) (uint32, error) {
-	for _, nalu := range nalus {
-		typ := h264.NALUType(nalu[0] & 0x1F)
-		if typ == h264.NALUTypeIDR || typ == h264.NALUTypeNonIDR {
-			poc, err := getPOC(nalu, sps)
-			if err != nil {
-				return 0, err
-			}
-			return poc, nil
-		}
-	}
-	return 0, fmt.Errorf("POC not found")
-}
-
-func getPOCDiff(poc uint32, expectedPOC uint32, sps *h264.SPS) int32 {
-	diff := int32(poc) - int32(expectedPOC)
-	switch {
-	case diff < -((1 << (sps.Log2MaxPicOrderCntLsbMinus4 + 3)) - 1):
-		diff += (1 << (sps.Log2MaxPicOrderCntLsbMinus4 + 4))
-
-	case diff > ((1 << (sps.Log2MaxPicOrderCntLsbMinus4 + 3)) - 1):
-		diff -= (1 << (sps.Log2MaxPicOrderCntLsbMinus4 + 4))
-	}
-	return diff
-}
-
 type fmp4VideoSample struct {
+	nalus      [][]byte
 	pts        time.Duration
 	dts        time.Duration
-	nalus      [][]byte
 	avcc       []byte
 	idrPresent bool
 	next       *fmp4VideoSample
-	pocDiff    int32
-}
-
-func (s *fmp4VideoSample) fillDTS(
-	prev *fmp4VideoSample,
-	sps *h264.SPS,
-	expectedPOC *uint32,
-) error {
-	if s.idrPresent || sps.PicOrderCntType == 2 {
-		s.dts = s.pts
-		*expectedPOC = 0
-	} else {
-		*expectedPOC += 2
-		*expectedPOC &= ((1 << (sps.Log2MaxPicOrderCntLsbMinus4 + 4)) - 1)
-
-		poc, err := getNALUSPOC(s.nalus, sps)
-		if err != nil {
-			return err
-		}
-
-		s.pocDiff = getPOCDiff(poc, *expectedPOC, sps)
-
-		if s.pocDiff == 0 {
-			s.dts = s.pts
-		} else {
-			if prev.pocDiff == 0 {
-				if s.pocDiff == -2 {
-					return fmt.Errorf("invalid frame POC")
-				}
-				s.dts = prev.pts + time.Duration(math.Round(float64(s.pts-prev.pts)/float64(s.pocDiff/2+1)))
-			} else {
-				s.dts = s.pts + time.Duration(math.Round(float64(prev.dts-prev.pts)*float64(s.pocDiff)/float64(prev.pocDiff)))
-			}
-		}
-	}
-
-	return nil
 }
 
 func (s fmp4VideoSample) duration() time.Duration {
@@ -183,8 +24,8 @@ func (s fmp4VideoSample) duration() time.Duration {
 }
 
 type fmp4AudioSample struct {
-	pts  time.Duration
 	au   []byte
+	pts  time.Duration
 	next *fmp4AudioSample
 }
 

internal/hls/muxer_variant_fmp4_segmenter.go

@@ -56,17 +56,14 @@ type muxerVariantFMP4Segmenter struct {
 
 	currentSegment        *muxerVariantFMP4Segment
 	startPTS              time.Duration
-	videoSPSP             *h264.SPS
 	videoSPS              []byte
 	videoPPS              []byte
-	videoNextSPSP         *h264.SPS
-	videoNextSPS          []byte
-	videoNextPPS          []byte
+	videoSPSP             *h264.SPS
+	videoDTSExtractor     *h264.DTSExtractor
 	nextSegmentID         uint64
 	nextPartID            uint64
 	nextVideoSample       *fmp4VideoSample
 	nextAudioSample       *fmp4AudioSample
-	videoExpectedPOC      uint32
 	firstSegmentFinalized bool
 	sampleDurations       map[time.Duration]struct{}
 	adjustedPartDuration  time.Duration
@@ -92,6 +89,7 @@ func newMuxerVariantFMP4Segmenter(
 		audioTrack:         audioTrack,
 		onSegmentFinalized: onSegmentFinalized,
 		onPartFinalized:    onPartFinalized,
+		videoDTSExtractor:  h264.NewDTSExtractor(),
 		nextSegmentID:      uint64(segmentCount),
 		sampleDurations:    make(map[time.Duration]struct{}),
 	}
@@ -109,16 +107,13 @@ func (m *muxerVariantFMP4Segmenter) genPartID() uint64 {
 	return id
 }
 
+// iPhone iOS fails if part durations are less than 85% of maximum part duration.
+// find a part duration that is compatible with all received sample durations
 func (m *muxerVariantFMP4Segmenter) adjustPartDuration(du time.Duration) {
-	if !m.lowLatency {
-		return
-	}
-	if m.firstSegmentFinalized {
+	if !m.lowLatency || m.firstSegmentFinalized {
 		return
 	}
 
-	// iPhone iOS fails if part durations are less than 85% of maximum part duration.
-	// find a part duration that is compatible with all received sample durations
 	if _, ok := m.sampleDurations[du]; !ok {
 		m.sampleDurations[du] = struct{}{}
 		m.adjustedPartDuration = findCompatiblePartDuration(
@@ -145,36 +140,44 @@ func (m *muxerVariantFMP4Segmenter) writeH264(pts time.Duration, nalus [][]byte)
 }
 
 func (m *muxerVariantFMP4Segmenter) writeH264Entry(sample *fmp4VideoSample) error {
-	// put SPS/PPS into a queue in order to sync them with the sample queue
-	m.videoSPSP = m.videoNextSPSP
-	m.videoSPS = m.videoNextSPS
-	m.videoPPS = m.videoNextPPS
+	// parse SPS
 	spsChanged := false
 	if sample.idrPresent {
-		videoNextSPS := m.videoTrack.SPS()
-		videoNextPPS := m.videoTrack.PPS()
+		videoSPS := m.videoTrack.SPS()
+		videoPPS := m.videoTrack.PPS()
 
 		if m.videoSPS == nil ||
-			!bytes.Equal(m.videoNextSPS, videoNextSPS) ||
-			!bytes.Equal(m.videoNextPPS, videoNextPPS) {
+			!bytes.Equal(m.videoSPS, videoSPS) ||
+			!bytes.Equal(m.videoPPS, videoPPS) {
 			spsChanged = true
 
 			var videoSPSP h264.SPS
-			err := videoSPSP.Unmarshal(videoNextSPS)
+			err := videoSPSP.Unmarshal(videoSPS)
 			if err != nil {
 				return err
 			}
 
-			m.videoNextSPSP = &videoSPSP
-			m.videoNextSPS = videoNextSPS
-			m.videoNextPPS = videoNextPPS
+			m.videoSPS = videoSPS
+			m.videoPPS = videoPPS
+			m.videoSPSP = &videoSPSP
 		}
 	}
 
+	// fill DTS
+	if m.videoSPSP != nil {
+		var err error
+		sample.dts, err = m.videoDTSExtractor.Extract(
+			sample.nalus, sample.idrPresent, sample.pts, m.videoSPSP)
+		if err != nil {
+			return err
+		}
+		sample.nalus = nil
+	}
+
 	sample.pts -= m.startPTS
+	sample.dts -= m.startPTS
 
 	// put samples into a queue in order to
-	// - allow to compute sample dts
 	// - allow to compute sample duration
 	// - check if next sample is IDR
 	sample, m.nextVideoSample = m.nextVideoSample, sample
@@ -206,18 +209,14 @@ func (m *muxerVariantFMP4Segmenter) writeH264Entry(sample *fmp4VideoSample) erro
 
 		m.startPTS = sample.pts
 		sample.pts = 0
+		sample.dts = 0
 		sample.next.pts -= m.startPTS
+		sample.next.dts -= m.startPTS
 	}
 
-	err := sample.next.fillDTS(sample, m.videoNextSPSP, &m.videoExpectedPOC)
-	if err != nil {
-		return err
-	}
-	sample.next.nalus = nil
-
 	m.adjustPartDuration(sample.duration())
 
-	err = m.currentSegment.writeH264(sample, m.adjustedPartDuration)
+	err := m.currentSegment.writeH264(sample, m.adjustedPartDuration)
 	if err != nil {
 		return err
 	}