diff --git a/example/test.go b/example/test.go index 96900fb..d86cdaa 100644 --- a/example/test.go +++ b/example/test.go @@ -211,15 +211,16 @@ func testInputGob(pathGob string, pathOut string, testSeg bool, writeM3u8 bool) writeM3U8Header(m3u8file) } - w := ts.SimpleH264Writer{ + muxer := &ts.Muxer{ W: outfile, - SPS: allSamples.SPS, - PPS: allSamples.PPS, - TimeScale: allSamples.TimeScale, } - lastPCR := int64(0) - //w.WriteHeader() + trackH264 := muxer.AddH264Track() + trackH264.SPS = allSamples.SPS + trackH264.PPS = allSamples.PPS + trackH264.TimeScale = int64(allSamples.TimeScale) + muxer.WriteHeader() + lastPTS := int64(0) syncCount := 0 segCount := 0 @@ -234,13 +235,13 @@ func testInputGob(pathGob string, pathOut string, testSeg bool, writeM3u8 bool) filename := fmt.Sprintf("%s.seg%d.ts", pathOut, segCount) if debugStream { - fmt.Println("stream:", "seg", segCount, "sync", syncCount, w.PCR) + fmt.Println("stream:", "seg", segCount, "sync", syncCount, trackH264.PTS) } if m3u8file != nil { info, _ := outfile.Stat() size := info.Size() - dur := float64(w.PCR - lastPCR) / float64(allSamples.TimeScale) + dur := float64(trackH264.PTS - lastPTS) / float64(allSamples.TimeScale) writeM3U8Item(m3u8file, lastFilename, size, dur) } @@ -248,14 +249,14 @@ func testInputGob(pathGob string, pathOut string, testSeg bool, writeM3u8 bool) outfile.Close() segCount++ outfile, _ = os.Create(filename) - w.W = outfile - w.WriteHeader() - lastPCR = w.PCR + muxer.W = outfile + muxer.WriteHeader() + lastPTS = trackH264.PTS } } - } - w.WriteNALU(sample.Sync, sample.Duration, sample.Data) + + trackH264.WriteH264NALU(sample.Sync, sample.Duration, sample.Data) } if m3u8file != nil { diff --git a/muxer.go b/muxer.go index aaa7017..fbec002 100644 --- a/muxer.go +++ b/muxer.go @@ -7,15 +7,46 @@ import ( ) type Track struct { - timeScale int64 - writeSPS bool SPS []byte PPS []byte - spsHasWritten bool - tsw *TSWriter + PTS int64 - PCR int64 - pesBuf *bytes.Buffer + TimeScale int64 + + writeSPS bool + spsHasWritten bool + pcrHasWritten bool + + streamId uint + tsw *TSWriter + dataBuf *iovec + cacheSize int +} + +func (self *Track) setPCR() { + if !self.pcrHasWritten { + self.tsw.PCR = 24300000 + self.pcrHasWritten = true + } else { + self.tsw.PCR = 0 + } +} + +func (self *Track) getPesHeader() (data []byte){ + if self.PTS == 0 { + self.PTS = self.TimeScale + } + buf := &bytes.Buffer{} + pes := PESHeader{ + StreamId: self.streamId, + PTS: uint64(self.PTS)*PTS_HZ/uint64(self.TimeScale), + } + WritePESHeader(buf, pes) + return buf.Bytes() +} + +func (self *Track) incPTS(delta int) { + self.PTS += int64(delta) } func (self *Track) WriteH264NALU(sync bool, duration int, nalu []byte) (err error) { @@ -28,16 +59,8 @@ func (self *Track) WriteH264NALU(sync bool, duration int, nalu []byte) (err erro } nalus = append(nalus, nalu) - pes := PESHeader{ - StreamId: StreamIdH264, - PTS: uint64(self.PTS)*PTS_HZ/uint64(self.timeScale), - } - if err = WritePESHeader(self.pesBuf, pes); err != nil { - return - } - data := &iovec{} - data.Append(self.pesBuf.Bytes()) + data.Append(self.getPesHeader()) for i, nalu := range nalus { var startCode []byte if i == 0 { @@ -50,64 +73,51 @@ func (self *Track) WriteH264NALU(sync bool, duration int, nalu []byte) (err erro } self.tsw.RandomAccessIndicator = sync - self.tsw.PCR = uint64(self.PCR)*PCR_HZ/uint64(self.timeScale) - + self.setPCR() if err = self.tsw.WriteIovec(data); err != nil { return } - self.PTS += int64(duration) - self.PCR += int64(duration) - self.pesBuf.Reset() - + self.incPTS(duration) return } func (self *Track) WriteADTSAACFrame(duration int, frame []byte) (err error) { - pes := PESHeader{ - StreamId: StreamIdAAC, - PTS: uint64(self.PTS)*PTS_HZ/uint64(self.timeScale), - } - if err = WritePESHeader(self.pesBuf, pes); err != nil { - return + if self.dataBuf != nil && self.dataBuf.Len > self.cacheSize { + self.tsw.RandomAccessIndicator = true + self.setPCR() + if err = self.tsw.WriteIovec(self.dataBuf); err != nil { + return + } + self.dataBuf = nil } - data := &iovec{} - data.Append(self.pesBuf.Bytes()) - data.Append(frame) - - self.tsw.RandomAccessIndicator = true - self.tsw.PCR = uint64(self.PCR)*PCR_HZ/uint64(self.timeScale) - if err = self.tsw.WriteIovec(data); err != nil { - return + if self.dataBuf == nil { + self.dataBuf = &iovec{} + self.dataBuf.Append(self.getPesHeader()) + } else { + self.dataBuf.Append(frame) } - self.PTS += int64(duration) - self.PCR += int64(duration) - self.pesBuf.Reset() - + self.incPTS(duration) return } -func newTrack(w io.Writer, pid uint, timeScale int64) (track *Track) { +func newTrack(w io.Writer, pid uint, streamId uint) (track *Track) { track = &Track{ tsw: &TSWriter{ W: w, PID: pid, - DiscontinuityIndicator: true, + //DiscontinuityIndicator: true, }, - timeScale: timeScale, - pesBuf: &bytes.Buffer{}, + streamId: streamId, } track.tsw.EnableVecWriter() - track.PTS = timeScale - track.PCR = timeScale return } type Muxer struct { W io.Writer - TimeScale int64 tswPAT *TSWriter tswPMT *TSWriter elemStreams []ElementaryStreamInfo @@ -118,7 +128,10 @@ func (self *Muxer) AddAACTrack() (track *Track) { self.elemStreams, ElementaryStreamInfo{StreamType: ElementaryStreamTypeAdtsAAC, ElementaryPID: 0x101}, ) - return newTrack(self.W, 0x101, self.TimeScale) + track = newTrack(self.W, 0x101, StreamIdAAC) + track.pcrHasWritten = true + track.cacheSize = 3000 + return } func (self *Muxer) AddH264Track() (track *Track) { @@ -126,7 +139,8 @@ func (self *Muxer) AddH264Track() (track *Track) { self.elemStreams, ElementaryStreamInfo{StreamType: ElementaryStreamTypeH264, ElementaryPID: 0x100}, ) - return newTrack(self.W, 0x100, self.TimeScale) + track = newTrack(self.W, 0x100, StreamIdH264) + return } func (self *Muxer) WriteHeader() (err error) { @@ -148,12 +162,12 @@ func (self *Muxer) WriteHeader() (err error) { tswPMT := &TSWriter{ W: self.W, PID: 0x1000, - DiscontinuityIndicator: true, + //DiscontinuityIndicator: true, } tswPAT := &TSWriter{ W: self.W, PID: 0, - DiscontinuityIndicator: true, + //DiscontinuityIndicator: true, } if err = tswPAT.Write(bufPAT.Bytes()); err != nil { return