h264: remove DTS estimator

examples/client-read-h264-save-to-disk/mpegtsencoder.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bufio"
+	"bytes"
 	"context"
 	"log"
 	"os"
@@ -11,6 +12,10 @@ import (
 	"github.com/asticode/go-astits"
 )
 
+const (
+	ptsDTSOffset = 400 * time.Millisecond
+)
+
 // mpegtsEncoder allows to encode H264 NALUs into MPEG-TS.
 type mpegtsEncoder struct {
 	sps []byte
@@ -19,9 +24,11 @@ type mpegtsEncoder struct {
 	f                  *os.File
 	b                  *bufio.Writer
 	mux                *astits.Muxer
-	dtsEst             *h264.DTSEstimator
+	dtsExtractor       *h264.DTSExtractor
 	firstPacketWritten bool
 	startPTS           time.Duration
+	spsp               *h264.SPS
+	firstIDRReceived   bool
 }
 
 // newMPEGTSEncoder allocates a mpegtsEncoder.
@@ -39,13 +46,23 @@ func newMPEGTSEncoder(sps []byte, pps []byte) (*mpegtsEncoder, error) {
 	})
 	mux.SetPCRPID(256)
 
+	var spsp *h264.SPS
+	if sps != nil {
+		spsp = &h264.SPS{}
+		err := spsp.Unmarshal(sps)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	return &mpegtsEncoder{
-		sps:    sps,
-		pps:    pps,
-		f:      f,
-		b:      b,
-		mux:    mux,
-		dtsEst: h264.NewDTSEstimator(),
+		sps:          sps,
+		pps:          pps,
+		f:            f,
+		b:            b,
+		mux:          mux,
+		dtsExtractor: h264.NewDTSExtractor(),
+		spsp:         spsp,
 	}, nil
 }
 
@@ -67,10 +84,20 @@ func (e *mpegtsEncoder) encode(nalus [][]byte, pts time.Duration) error {
 		{byte(h264.NALUTypeAccessUnitDelimiter), 240},
 	}
 
+	idrPresent := h264.IDRPresent(nalus)
+
 	for _, nalu := range nalus {
 		typ := h264.NALUType(nalu[0] & 0x1F)
 		switch typ {
 		case h264.NALUTypeSPS:
+			if e.sps == nil || !bytes.Equal(e.sps, nalu) {
+				var spsp h264.SPS
+				err := spsp.Unmarshal(nalu)
+				if err != nil {
+					return err
+				}
+				e.spsp = &spsp
+			}
 			e.sps = append([]byte(nil), nalu...)
 			continue
 
@@ -91,11 +118,15 @@ func (e *mpegtsEncoder) encode(nalus [][]byte, pts time.Duration) error {
 		filteredNALUs = append(filteredNALUs, nalu)
 	}
 
-	// it's useless to go on if SPS or PPS have not been provided yet
 	if e.sps == nil || e.pps == nil {
 		return nil
 	}
 
+	if !e.firstIDRReceived && !idrPresent {
+		return nil
+	}
+	e.firstIDRReceived = true
+
 	// encode into Annex-B
 	enc, err := h264.AnnexBEncode(filteredNALUs)
 	if err != nil {
@@ -103,7 +134,13 @@ func (e *mpegtsEncoder) encode(nalus [][]byte, pts time.Duration) error {
 	}
 
 	pts -= e.startPTS
-	dts := e.dtsEst.Feed(pts)
+
+	dts, err := e.dtsExtractor.Extract(filteredNALUs, idrPresent, pts, e.spsp)
+	if err != nil {
+		return err
+	}
+
+	pts += ptsDTSOffset
 
 	oh := &astits.PESOptionalHeader{
 		MarkerBits: 2,

pkg/h264/dtsestimator.go

@@ -1,51 +0,0 @@
-package h264
-
-import (
-	"time"
-)
-
-// DTSEstimator is a DTS estimator.
-type DTSEstimator struct {
-	initializing int
-	prevDTS      time.Duration
-	prevPTS      time.Duration
-}
-
-// NewDTSEstimator allocates a DTSEstimator.
-func NewDTSEstimator() *DTSEstimator {
-	return &DTSEstimator{
-		initializing: 2,
-	}
-}
-
-// Feed provides PTS to the estimator, and returns the estimated DTS.
-func (d *DTSEstimator) Feed(pts time.Duration) time.Duration {
-	switch d.initializing {
-	case 2:
-		d.initializing--
-		return 0
-
-	case 1:
-		d.initializing--
-		d.prevPTS = pts
-		d.prevDTS = time.Millisecond
-		return time.Millisecond
-	}
-
-	dts := func() time.Duration {
-		// PTS is increasing
-		// use previous PTS
-		if pts > d.prevPTS {
-			return d.prevPTS
-		}
-
-		// PTS is not increasing
-		// use last DTS value plus a small quantity
-		return d.prevDTS + time.Millisecond
-	}()
-
-	d.prevPTS = pts
-	d.prevDTS = dts
-
-	return dts
-}

pkg/h264/dtsestimator_test.go

@@ -1,32 +0,0 @@
-package h264
-
-import (
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/require"
-)
-
-func TestDTSEstimator(t *testing.T) {
-	est := NewDTSEstimator()
-
-	// initial state
-	dts := est.Feed(0)
-	require.Equal(t, time.Duration(0), dts)
-
-	// B frame
-	dts = est.Feed(1*time.Second - 200*time.Millisecond)
-	require.Equal(t, time.Millisecond, dts)
-
-	// B frame
-	dts = est.Feed(1*time.Second - 400*time.Millisecond)
-	require.Equal(t, 2*time.Millisecond, dts)
-
-	// P frame
-	dts = est.Feed(1 * time.Second)
-	require.Equal(t, 1*time.Second-400*time.Millisecond, dts)
-
-	// P frame
-	dts = est.Feed(1*time.Second + 200*time.Millisecond)
-	require.Equal(t, 1*time.Second, dts)
-}