diff --git a/demuxer.go b/demuxer.go index b6b18b9..97bb4cc 100644 --- a/demuxer.go +++ b/demuxer.go @@ -138,10 +138,11 @@ func (self *Track) ReadSample() (pts int64, dts int64, isKeyFrame bool, data []b } if self.Type == AAC { - var n int - if _, data, n, self.payload, err = isom.ReadADTSPayload(self.payload); err != nil { + var n, framelen int + if _, data, n, framelen, err = isom.ReadADTSFrame(self.payload); err != nil { return } + self.payload = self.payload[framelen:] pts = self.PTS dts = pts self.PTS += int64(PTS_HZ*n)/int64(self.mpeg4AudioConfig.SampleRate) @@ -151,14 +152,14 @@ func (self *Track) ReadSample() (pts int64, dts int64, isKeyFrame bool, data []b } else { dts = int64(self.peshdr.DTS) pts = int64(self.peshdr.PTS) + if dts == 0 { + dts = pts + } isKeyFrame = self.tshdr.RandomAccessIndicator data = self.payload self.payloadReady = false } - if dts == 0 { - dts = pts - } return } @@ -171,7 +172,7 @@ func (self *Track) appendPayload() (err error) { if self.Type == AAC { if !self.mpeg4AudioConfig.IsValid() { - if self.mpeg4AudioConfig, _, _, _, err = isom.ReadADTSPayload(self.payload); err != nil { + if self.mpeg4AudioConfig, _, _, _, err = isom.ReadADTSFrame(self.payload); err != nil { return } self.mpeg4AudioConfig = self.mpeg4AudioConfig.Complete() diff --git a/muxer.go b/muxer.go index f8a634b..b13d4a0 100644 --- a/muxer.go +++ b/muxer.go @@ -4,20 +4,184 @@ package ts import ( "bytes" "io" + "github.com/nareix/mp4/isom" ) +type Muxer struct { + W io.Writer + tswPAT *TSWriter + tswPMT *TSWriter + elemStreams []ElementaryStreamInfo + TrackH264 *Track + Tracks []*Track +} + +func (self *Muxer) newTrack(pid uint, streamId uint) (track *Track) { + track = &Track{ + mux: self, + tsw: &TSWriter{ + PID: pid, + DiscontinuityIndicator: true, + }, + streamId: streamId, + } + track.tsw.EnableVecWriter() + return +} + +func (self *Muxer) AddAACTrack() (track *Track) { + self.elemStreams = append( + self.elemStreams, + ElementaryStreamInfo{StreamType: ElementaryStreamTypeAdtsAAC, ElementaryPID: 0x101}, + ) + track = self.newTrack(0x101, StreamIdAAC) + track.Type = AAC + track.cacheSize = 3000 + self.Tracks = append(self.Tracks, track) + return +} + +func (self *Muxer) AddH264Track() (track *Track) { + self.elemStreams = append( + self.elemStreams, + ElementaryStreamInfo{StreamType: ElementaryStreamTypeH264, ElementaryPID: 0x100}, + ) + track = self.newTrack(0x100, StreamIdH264) + track.Type = H264 + self.TrackH264 = track + self.Tracks = append(self.Tracks, track) + return +} + +func (self *Muxer) WriteHeader() (err error) { + bufPAT := &bytes.Buffer{} + bufPMT := &bytes.Buffer{} + + pat := PAT{ + Entries: []PATEntry{ + {ProgramNumber: 1, ProgramMapPID: 0x1000}, + }, + } + WritePAT(bufPAT, pat) + pmt := PMT{ + PCRPID: 0x100, + ElementaryStreamInfos: self.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 + } + + for _, track := range(self.Tracks) { + track.spsHasWritten = false + } + + return +} + +func (self *Track) SetH264PPSAndSPS(pps []byte, sps []byte) { + self.PPS, self.SPS = pps, sps +} + +func (self *Track) SetTimeScale(timeScale int64) { + self.timeScale = timeScale +} + +func (self *Track) TimeScale() int64 { + return self.timeScale +} + +func (self *Track) SetMPEG4AudioConfig(config isom.MPEG4AudioConfig) { + self.mpeg4AudioConfig = config +} + +func (self *Track) tsToPesTs(ts int64) uint64 { + return uint64(ts)*PTS_HZ/uint64(self.timeScale) +} + +func (self *Track) tsToPCR(ts int64) uint64 { + return uint64(ts)*PCR_HZ/uint64(self.timeScale)+PCR_HZ +} + +func (self *Track) WriteSample(pts int64, dts int64, isKeyFrame bool, data []byte) (err error) { + if self.Type == AAC { + + if !isom.IsADTSFrame(data) { + data = append(data, isom.MakeADTSHeader(self.mpeg4AudioConfig, 1024, len(data))...) + } + + buf := &bytes.Buffer{} + pes := PESHeader{ + StreamId: self.streamId, + PTS: self.tsToPesTs(pts), + } + WritePESHeader(buf, pes, len(data)) + buf.Write(data) + + self.tsw.RandomAccessIndicator = true + self.tsw.PCR = self.tsToPCR(dts) + if err = self.tsw.WriteTo(self.mux.W, buf.Bytes()); err != nil { + return + } + } else if self.Type == H264 { + + buf := &bytes.Buffer{} + pes := PESHeader{ + StreamId: self.streamId, + PTS: self.tsToPesTs(pts), + } + if dts != pts { + pes.DTS = self.tsToPesTs(dts) + } + WritePESHeader(buf, pes, 0) + + if isKeyFrame { + buf.Write([]byte{0,0,0,1,0x9,0xf0,0,0,0,1}) // AUD + buf.Write(self.SPS) + buf.Write([]byte{0,0,1}) + buf.Write(self.PPS) + buf.Write([]byte{0,0,1}) + buf.Write(data) + } else { + buf.Write([]byte{0,0,0,1,0x9,0xf0,0,0,0,1}) // AUD + buf.Write(data) + } + + self.tsw.RandomAccessIndicator = isKeyFrame + self.tsw.PCR = self.tsToPCR(dts) + if err = self.tsw.WriteTo(self.mux.W, buf.Bytes()); err != nil { + return + } + } + return +} + +/* about to remove */ + func (self *Track) setPCR() { - self.tsw.PCR = uint64(self.PTS)*PCR_HZ/uint64(self.TimeScale) + self.tsw.PCR = uint64(self.PTS)*PCR_HZ/uint64(self.timeScale) } func (self *Track) getPesHeader(dataLength int) (data []byte){ if self.PTS == 0 { - self.PTS = self.TimeScale + self.PTS = self.timeScale } buf := &bytes.Buffer{} pes := PESHeader{ StreamId: self.streamId, - PTS: uint64(self.PTS)*PTS_HZ/uint64(self.TimeScale), + PTS: uint64(self.PTS)*PTS_HZ/uint64(self.timeScale), } WritePESHeader(buf, pes, dataLength) return buf.Bytes() @@ -78,85 +242,3 @@ func (self *Track) WriteADTSAACFrame(duration int, frame []byte) (err error) { return } -func newTrack(mux *Muxer, pid uint, streamId uint) (track *Track) { - track = &Track{ - mux: mux, - tsw: &TSWriter{ - PID: pid, - DiscontinuityIndicator: true, - }, - streamId: streamId, - } - track.tsw.EnableVecWriter() - return -} - -type Muxer struct { - W io.Writer - tswPAT *TSWriter - tswPMT *TSWriter - elemStreams []ElementaryStreamInfo - TrackH264 *Track - Tracks []*Track -} - -func (self *Muxer) AddAACTrack() (track *Track) { - self.elemStreams = append( - self.elemStreams, - ElementaryStreamInfo{StreamType: ElementaryStreamTypeAdtsAAC, ElementaryPID: 0x101}, - ) - track = newTrack(self, 0x101, StreamIdAAC) - track.cacheSize = 3000 - self.Tracks = append(self.Tracks, track) - return -} - -func (self *Muxer) AddH264Track() (track *Track) { - self.elemStreams = append( - self.elemStreams, - ElementaryStreamInfo{StreamType: ElementaryStreamTypeH264, ElementaryPID: 0x100}, - ) - track = newTrack(self, 0x100, StreamIdH264) - self.TrackH264 = track - self.Tracks = append(self.Tracks, track) - return -} - -func (self *Muxer) WriteHeader() (err error) { - bufPAT := &bytes.Buffer{} - bufPMT := &bytes.Buffer{} - - pat := PAT{ - Entries: []PATEntry{ - {ProgramNumber: 1, ProgramMapPID: 0x1000}, - }, - } - WritePAT(bufPAT, pat) - pmt := PMT{ - PCRPID: 0x100, - ElementaryStreamInfos: self.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 - } - - for _, track := range(self.Tracks) { - track.spsHasWritten = false - } - - return -} - diff --git a/track.go b/track.go index a83aa4c..cb407cf 100644 --- a/track.go +++ b/track.go @@ -14,7 +14,7 @@ type Track struct { pid uint PTS int64 - TimeScale int64 + timeScale int64 mpeg4AudioConfig isom.MPEG4AudioConfig buf bytes.Buffer