hls: move NALU filtering into hls muxer

This commit is contained in:
aler9 2021-08-14 15:39:29 +02:00
parent d5b10fdd66
commit 600f7bf48c
4 changed files with 62 additions and 49 deletions

View File

@ -18,7 +18,6 @@ import (
"github.com/aler9/gortsplib/pkg/rtph264"
"github.com/pion/rtp"
"github.com/aler9/rtsp-simple-server/internal/h264"
"github.com/aler9/rtsp-simple-server/internal/hls"
"github.com/aler9/rtsp-simple-server/internal/logger"
)
@ -244,8 +243,6 @@ func (r *hlsRemuxer) runRemuxer(remuxerCtx context.Context, remuxerReady chan st
var videoTrack *gortsplib.Track
videoTrackID := -1
var h264SPS []byte
var h264PPS []byte
var h264Decoder *rtph264.Decoder
var audioTrack *gortsplib.Track
audioTrackID := -1
@ -261,12 +258,6 @@ func (r *hlsRemuxer) runRemuxer(remuxerCtx context.Context, remuxerReady chan st
videoTrack = t
videoTrackID = i
var err error
h264SPS, h264PPS, err = t.ExtractDataH264()
if err != nil {
return err
}
h264Decoder = rtph264.NewDecoder()
} else if t.IsAAC() {
@ -341,22 +332,7 @@ func (r *hlsRemuxer) runRemuxer(remuxerCtx context.Context, remuxerReady chan st
continue
}
for _, nalu := range nalus {
// remove SPS, PPS, AUD
typ := h264.NALUType(nalu[0] & 0x1F)
switch typ {
case h264.NALUTypeSPS, h264.NALUTypePPS, h264.NALUTypeAccessUnitDelimiter:
continue
}
// add SPS and PPS before IDR
if typ == h264.NALUTypeIDR {
videoBuf = append(videoBuf, h264SPS)
videoBuf = append(videoBuf, h264PPS)
}
videoBuf = append(videoBuf, nalu)
}
videoBuf = append(videoBuf, nalus...)
// RTP marker means that all the NALUs with the same PTS have been received.
// send them together.

View File

@ -31,6 +31,8 @@ type Muxer struct {
videoTrack *gortsplib.Track
audioTrack *gortsplib.Track
h264SPS []byte
h264PPS []byte
aacConfig rtpaac.MPEG4AudioConfig
startPCR time.Time
videoDTSEst *h264.DTSEstimator
@ -48,6 +50,16 @@ func NewMuxer(
hlsSegmentDuration time.Duration,
videoTrack *gortsplib.Track,
audioTrack *gortsplib.Track) (*Muxer, error) {
var h264SPS []byte
var h264PPS []byte
if videoTrack != nil {
var err error
h264SPS, h264PPS, err = videoTrack.ExtractDataH264()
if err != nil {
return nil, err
}
}
var aacConfig rtpaac.MPEG4AudioConfig
if audioTrack != nil {
byts, err := audioTrack.ExtractDataAAC()
@ -66,10 +78,12 @@ func NewMuxer(
hlsSegmentDuration: hlsSegmentDuration,
videoTrack: videoTrack,
audioTrack: audioTrack,
h264SPS: h264SPS,
h264PPS: h264PPS,
aacConfig: aacConfig,
startPCR: time.Now(),
videoDTSEst: h264.NewDTSEstimator(),
tsCurrent: newTSFile(videoTrack != nil, audioTrack != nil),
tsCurrent: newTSFile(videoTrack, audioTrack),
tsByName: make(map[string]*tsFile),
}
@ -111,7 +125,7 @@ func (m *Muxer) WriteH264(pts time.Duration, nalus [][]byte) error {
m.tsCurrent.close()
}
m.tsCurrent = newTSFile(m.videoTrack != nil, m.audioTrack != nil)
m.tsCurrent = newTSFile(m.videoTrack, m.audioTrack)
m.tsByName[m.tsCurrent.name] = m.tsCurrent
m.tsQueue = append(m.tsQueue, m.tsCurrent)
@ -124,6 +138,8 @@ func (m *Muxer) WriteH264(pts time.Duration, nalus [][]byte) error {
m.tsCurrent.setPCR(time.Since(m.startPCR))
err := m.tsCurrent.writeH264(
m.h264SPS,
m.h264PPS,
m.videoDTSEst.Feed(pts+ptsOffset),
pts+ptsOffset,
idrPresent,
@ -150,7 +166,7 @@ func (m *Muxer) WriteAAC(pts time.Duration, aus [][]byte) error {
}
m.audioAUCount = 0
m.tsCurrent = newTSFile(m.videoTrack != nil, m.audioTrack != nil)
m.tsCurrent = newTSFile(m.videoTrack, m.audioTrack)
m.tsByName[m.tsCurrent.name] = m.tsCurrent
m.tsQueue = append(m.tsQueue, m.tsCurrent)
if len(m.tsQueue) > m.hlsSegmentCount {

View File

@ -17,7 +17,7 @@ func checkTSPacket(t *testing.T, byts []byte, pid int, afc int) {
}
func TestMuxer(t *testing.T) {
videoTrack, err := gortsplib.NewTrackH264(96, []byte{0x01, 0x02, 0x03, 0x04}, []byte{0x01, 0x02, 0x03, 0x04})
videoTrack, err := gortsplib.NewTrackH264(96, []byte{0x07, 0x01, 0x02, 0x03}, []byte{0x08})
require.NoError(t, err)
audioTrack, err := gortsplib.NewTrackAAC(97, []byte{17, 144})
@ -92,12 +92,11 @@ func TestMuxer(t *testing.T) {
byts = byts[4+alen+20:]
require.Equal(t,
[]byte{
0, 0, 0, 1, 9, 240,
0, 0, 0, 1, 5,
0, 0, 0, 1, 9,
0, 0, 0, 1, 8,
0, 0, 0, 1, 7,
0, 0, 0, 1, 9, 240, // AUD
0, 0, 0, 1, 7, 1, 2, 3, // SPS
0, 0, 0, 1, 8, // PPS
0, 0, 0, 1, 5, // IDR
},
byts[:26],
byts[:24],
)
}

View File

@ -6,6 +6,7 @@ import (
"strconv"
"time"
"github.com/aler9/gortsplib"
"github.com/asticode/go-astits"
"github.com/aler9/rtsp-simple-server/internal/aac"
@ -13,7 +14,7 @@ import (
)
type tsFile struct {
hasVideoTrack bool
videoTrack *gortsplib.Track
name string
buf *multiAccessBuffer
mux *astits.Muxer
@ -23,30 +24,30 @@ type tsFile struct {
maxPTS time.Duration
}
func newTSFile(hasVideoTrack bool, hasAudioTrack bool) *tsFile {
func newTSFile(videoTrack *gortsplib.Track, audioTrack *gortsplib.Track) *tsFile {
t := &tsFile{
hasVideoTrack: hasVideoTrack,
name: strconv.FormatInt(time.Now().Unix(), 10),
buf: newMultiAccessBuffer(),
videoTrack: videoTrack,
name: strconv.FormatInt(time.Now().Unix(), 10),
buf: newMultiAccessBuffer(),
}
t.mux = astits.NewMuxer(context.Background(), t.buf)
if hasVideoTrack {
if videoTrack != nil {
t.mux.AddElementaryStream(astits.PMTElementaryStream{
ElementaryPID: 256,
StreamType: astits.StreamTypeH264Video,
})
}
if hasAudioTrack {
if audioTrack != nil {
t.mux.AddElementaryStream(astits.PMTElementaryStream{
ElementaryPID: 257,
StreamType: astits.StreamTypeAACAudio,
})
}
if hasVideoTrack {
if videoTrack != nil {
t.mux.SetPCRPID(256)
} else {
t.mux.SetPCRPID(257)
@ -76,7 +77,9 @@ func (t *tsFile) newReader() io.Reader {
return t.buf.NewReader()
}
func (t *tsFile) writeH264(dts time.Duration, pts time.Duration, isIDR bool, nalus [][]byte) error {
func (t *tsFile) writeH264(
h264SPS []byte, h264PPS []byte,
dts time.Duration, pts time.Duration, isIDR bool, nalus [][]byte) error {
if !t.firstPacketWritten {
t.firstPacketWritten = true
t.minPTS = pts
@ -90,10 +93,29 @@ func (t *tsFile) writeH264(dts time.Duration, pts time.Duration, isIDR bool, nal
}
}
// prepend an AUD. This is required by video.js and iOS
nalus = append([][]byte{{byte(h264.NALUTypeAccessUnitDelimiter), 240}}, nalus...)
filteredNALUs := [][]byte{
// prepend an AUD. This is required by video.js and iOS
{byte(h264.NALUTypeAccessUnitDelimiter), 240},
}
enc, err := h264.EncodeAnnexB(nalus)
for _, nalu := range nalus {
// remove existing SPS, PPS, AUD
typ := h264.NALUType(nalu[0] & 0x1F)
switch typ {
case h264.NALUTypeSPS, h264.NALUTypePPS, h264.NALUTypeAccessUnitDelimiter:
continue
}
// add SPS and PPS before IDR
if typ == h264.NALUTypeIDR {
filteredNALUs = append(filteredNALUs, h264SPS)
filteredNALUs = append(filteredNALUs, h264PPS)
}
filteredNALUs = append(filteredNALUs, nalu)
}
enc, err := h264.EncodeAnnexB(filteredNALUs)
if err != nil {
return err
}
@ -124,7 +146,7 @@ func (t *tsFile) writeH264(dts time.Duration, pts time.Duration, isIDR bool, nal
}
func (t *tsFile) writeAAC(sampleRate int, channelCount int, pts time.Duration, au []byte) error {
if !t.hasVideoTrack {
if t.videoTrack == nil {
if !t.firstPacketWritten {
t.firstPacketWritten = true
t.minPTS = pts
@ -154,7 +176,7 @@ func (t *tsFile) writeAAC(sampleRate int, channelCount int, pts time.Duration, a
RandomAccessIndicator: true,
}
if !t.hasVideoTrack {
if t.videoTrack == nil {
af.HasPCR = true
af.PCR = &astits.ClockReference{Base: int64(t.pcr.Seconds() * 90000)}
}