135 lines
2.9 KiB
Go
135 lines
2.9 KiB
Go
package ts
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"github.com/nareix/av"
|
|
"github.com/nareix/codec/aacparser"
|
|
"github.com/nareix/codec/h264parser"
|
|
"io"
|
|
)
|
|
|
|
type Muxer struct {
|
|
W io.Writer
|
|
streams []*Stream
|
|
}
|
|
|
|
func (self *Muxer) NewStream() av.Stream {
|
|
stream := &Stream{
|
|
mux: self,
|
|
tsw: &TSWriter{
|
|
DiscontinuityIndicator: true,
|
|
PID: uint(len(self.streams) + 0x100),
|
|
},
|
|
}
|
|
self.streams = append(self.streams, stream)
|
|
return stream
|
|
}
|
|
|
|
func (self *Muxer) WriteHeader() (err error) {
|
|
bufPAT := &bytes.Buffer{}
|
|
bufPMT := &bytes.Buffer{}
|
|
|
|
pat := PAT{
|
|
Entries: []PATEntry{
|
|
{ProgramNumber: 1, ProgramMapPID: 0x1000},
|
|
},
|
|
}
|
|
WritePAT(bufPAT, pat)
|
|
|
|
var elemStreams []ElementaryStreamInfo
|
|
for _, stream := range self.streams {
|
|
switch stream.Type() {
|
|
case av.AAC:
|
|
elemStreams = append(elemStreams, ElementaryStreamInfo{StreamType: ElementaryStreamTypeAdtsAAC, ElementaryPID: stream.tsw.PID})
|
|
case av.H264:
|
|
elemStreams = append(elemStreams, ElementaryStreamInfo{StreamType: ElementaryStreamTypeH264, ElementaryPID: stream.tsw.PID})
|
|
}
|
|
}
|
|
|
|
pmt := PMT{
|
|
PCRPID: 0x100,
|
|
ElementaryStreamInfos: elemStreams,
|
|
}
|
|
WritePMT(bufPMT, pmt)
|
|
|
|
tswPMT := &TSWriter{
|
|
PID: 0x1000,
|
|
DiscontinuityIndicator: true,
|
|
}
|
|
tswPAT := &TSWriter{
|
|
PID: 0,
|
|
DiscontinuityIndicator: true,
|
|
}
|
|
if err = tswPAT.WriteTo(self.W, bufPAT.Bytes()); err != nil {
|
|
return
|
|
}
|
|
if err = tswPMT.WriteTo(self.W, bufPMT.Bytes()); err != nil {
|
|
return
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
func (self *Muxer) WritePacket(streamIndex int, pkt av.Packet) (err error) {
|
|
stream := self.streams[streamIndex]
|
|
|
|
if stream.Type() == av.AAC {
|
|
data := pkt.Data
|
|
if !aacparser.IsADTSFrame(data) {
|
|
data = append(aacparser.MakeADTSHeader(stream.AACCodecInfo.MPEG4AudioConfig, 1024, len(data)), data...)
|
|
}
|
|
|
|
buf := &bytes.Buffer{}
|
|
pes := PESHeader{
|
|
StreamId: StreamIdAAC,
|
|
PTS: timeToPesTs(stream.time),
|
|
}
|
|
WritePESHeader(buf, pes, len(data))
|
|
buf.Write(data)
|
|
|
|
stream.tsw.RandomAccessIndicator = true
|
|
stream.tsw.PCR = timeToPCR(stream.time)
|
|
if err = stream.tsw.WriteTo(self.W, buf.Bytes()); err != nil {
|
|
return
|
|
}
|
|
|
|
stream.time += pkt.Duration
|
|
|
|
} else if stream.Type() == av.H264 {
|
|
buf := &bytes.Buffer{}
|
|
pes := PESHeader{
|
|
StreamId: StreamIdH264,
|
|
PTS: timeToPesTs(stream.time),
|
|
}
|
|
if pkt.CompositionTime > 0.0 {
|
|
pes.DTS = timeToPesTs(stream.time + pkt.CompositionTime)
|
|
}
|
|
WritePESHeader(buf, pes, 0)
|
|
|
|
nalus, _ := h264parser.SplitNALUs(pkt.Data)
|
|
if pkt.IsKeyFrame {
|
|
sps := stream.H264CodecInfo.Record.SPS[0]
|
|
pps := stream.H264CodecInfo.Record.PPS[0]
|
|
nalus = append([][]byte{sps, pps}, nalus...)
|
|
}
|
|
h264parser.WalkNALUsAnnexb(nalus, func(b []byte) {
|
|
buf.Write(b)
|
|
})
|
|
|
|
stream.tsw.RandomAccessIndicator = pkt.IsKeyFrame
|
|
stream.tsw.PCR = timeToPCR(stream.time)
|
|
if err = stream.tsw.WriteTo(self.W, buf.Bytes()); err != nil {
|
|
return
|
|
}
|
|
|
|
stream.time += pkt.Duration
|
|
|
|
} else {
|
|
err = fmt.Errorf("unknown stream type=%d", stream.Type())
|
|
return
|
|
}
|
|
|
|
return
|
|
}
|