diff --git a/examples/client-read-h264-save-to-disk/main.go b/examples/client-read-h264-save-to-disk/main.go index e1aa1e11..e5d3f227 100644 --- a/examples/client-read-h264-save-to-disk/main.go +++ b/examples/client-read-h264-save-to-disk/main.go @@ -1,18 +1,10 @@ package main import ( - "bufio" - "context" - "fmt" - "os" - "time" - "github.com/aler9/gortsplib" "github.com/aler9/gortsplib/pkg/base" - "github.com/aler9/gortsplib/pkg/h264" "github.com/aler9/gortsplib/pkg/headers" "github.com/aler9/gortsplib/pkg/rtph264" - "github.com/asticode/go-astits" "github.com/pion/rtp" ) @@ -22,30 +14,9 @@ import ( // 3. save the content of the H264 track to a file in MPEG-TS format func main() { - // open output file - f, err := os.Create("mystream.ts") - if err != nil { - panic(err) - } - defer f.Close() - - // istantiate things needed to decode RTP/H264 and encode MPEG-TS - b := bufio.NewWriter(f) - defer b.Flush() - mux := astits.NewMuxer(context.Background(), b) dec := rtph264.NewDecoder() - dtsEst := h264.NewDTSEstimator() - firstPacketWritten := false - var startPTS time.Duration var h264Track int - var h264Conf *gortsplib.TrackConfigH264 - - // add an H264 track to the MPEG-TS muxer - mux.AddElementaryStream(astits.PMTElementaryStream{ - ElementaryPID: 256, - StreamType: astits.StreamTypeH264Video, - }) - mux.SetPCRPID(256) + var enc *mpegtsEncoder c := gortsplib.Client{ // called when a RTP packet arrives @@ -54,7 +25,7 @@ func main() { return } - // parse RTP packets + // parse RTP packet var pkt rtp.Packet err := pkt.Unmarshal(payload) if err != nil { @@ -67,77 +38,11 @@ func main() { return } - if !firstPacketWritten { - firstPacketWritten = true - startPTS = pts - } - - // check whether there's an IDR - idrPresent := func() bool { - for _, nalu := range nalus { - typ := h264.NALUType(nalu[0] & 0x1F) - if typ == h264.NALUTypeIDR { - return true - } - } - return false - }() - - // prepend an AUD. This is required by some players - filteredNALUs := [][]byte{ - {byte(h264.NALUTypeAccessUnitDelimiter), 240}, - } - - 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 every IDR - if typ == h264.NALUTypeIDR { - filteredNALUs = append(filteredNALUs, h264Conf.SPS) - filteredNALUs = append(filteredNALUs, h264Conf.PPS) - } - - filteredNALUs = append(filteredNALUs, nalu) - } - - // encode into Annex-B - enc, err := h264.EncodeAnnexB(filteredNALUs) + // encode H264 NALUs into MPEG-TS + err = enc.encode(nalus, pts) if err != nil { - panic(err) + return } - - dts := dtsEst.Feed(pts - startPTS) - pts = pts - startPTS - - // write TS packet - _, err = mux.WriteData(&astits.MuxerData{ - PID: 256, - AdaptationField: &astits.PacketAdaptationField{ - RandomAccessIndicator: idrPresent, - }, - PES: &astits.PESData{ - Header: &astits.PESHeader{ - OptionalHeader: &astits.PESOptionalHeader{ - MarkerBits: 2, - PTSDTSIndicator: astits.PTSDTSIndicatorBothPresent, - DTS: &astits.ClockReference{Base: int64(dts.Seconds() * 90000)}, - PTS: &astits.ClockReference{Base: int64(pts.Seconds() * 90000)}, - }, - StreamID: 224, // video - }, - Data: enc, - }, - }) - if err != nil { - panic(err) - } - - fmt.Println("wrote ts packet") }, } @@ -175,18 +80,20 @@ func main() { return -1 }() if h264Track < 0 { - panic(fmt.Errorf("H264 track not found")) + panic("H264 track not found") } - fmt.Printf("H264 track is number %d\n", h264Track+1) // get track config - h264Conf, err = c.Tracks()[h264Track].ExtractConfigH264() + h264Conf, err := tracks[h264Track].ExtractConfigH264() if err != nil { panic(err) } - // instantiate a RTP/H264 decoder - dec = rtph264.NewDecoder() + // setup the encoder + enc, err = newMPEGTSEncoder(h264Conf) + if err != nil { + panic(err) + } // setup all tracks for _, t := range tracks { diff --git a/examples/client-read-h264-save-to-disk/mpegtsencoder.go b/examples/client-read-h264-save-to-disk/mpegtsencoder.go new file mode 100644 index 00000000..6c1c6bda --- /dev/null +++ b/examples/client-read-h264-save-to-disk/mpegtsencoder.go @@ -0,0 +1,130 @@ +package main + +import ( + "bufio" + "context" + "fmt" + "os" + "time" + + "github.com/aler9/gortsplib" + "github.com/aler9/gortsplib/pkg/h264" + "github.com/asticode/go-astits" +) + +// mpegtsEncoder allows to encode H264 NALUs into MPEG-TS. +type mpegtsEncoder struct { + f *os.File + b *bufio.Writer + mux *astits.Muxer + dtsEst *h264.DTSEstimator + firstPacketWritten bool + startPTS time.Duration + h264Conf *gortsplib.TrackConfigH264 +} + +// newMPEGTSEncoder allocates a mpegtsEncoder. +func newMPEGTSEncoder(h264Conf *gortsplib.TrackConfigH264) (*mpegtsEncoder, error) { + f, err := os.Create("mystream.ts") + if err != nil { + return nil, err + } + b := bufio.NewWriter(f) + + mux := astits.NewMuxer(context.Background(), b) + mux.AddElementaryStream(astits.PMTElementaryStream{ + ElementaryPID: 256, + StreamType: astits.StreamTypeH264Video, + }) + mux.SetPCRPID(256) + + return &mpegtsEncoder{ + f: f, + b: b, + mux: mux, + dtsEst: h264.NewDTSEstimator(), + h264Conf: h264Conf, + }, nil +} + +// close closes all the mpegtsEncoder resources. +func (e *mpegtsEncoder) close() { + e.b.Flush() + e.f.Close() +} + +// encode encodes H264 NALUs into MPEG-TS. +func (e *mpegtsEncoder) encode(nalus [][]byte, pts time.Duration) error { + if !e.firstPacketWritten { + e.firstPacketWritten = true + e.startPTS = pts + } + + // check whether there's an IDR + idrPresent := func() bool { + for _, nalu := range nalus { + typ := h264.NALUType(nalu[0] & 0x1F) + if typ == h264.NALUTypeIDR { + return true + } + } + return false + }() + + // prepend an AUD. This is required by some players + filteredNALUs := [][]byte{ + {byte(h264.NALUTypeAccessUnitDelimiter), 240}, + } + + 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 every IDR + if typ == h264.NALUTypeIDR { + filteredNALUs = append(filteredNALUs, e.h264Conf.SPS) + filteredNALUs = append(filteredNALUs, e.h264Conf.PPS) + } + + filteredNALUs = append(filteredNALUs, nalu) + } + + // encode into Annex-B + enc, err := h264.EncodeAnnexB(filteredNALUs) + if err != nil { + return err + } + + dts := e.dtsEst.Feed(pts - e.startPTS) + pts = pts - e.startPTS + + // write TS packet + _, err = e.mux.WriteData(&astits.MuxerData{ + PID: 256, + AdaptationField: &astits.PacketAdaptationField{ + RandomAccessIndicator: idrPresent, + }, + PES: &astits.PESData{ + Header: &astits.PESHeader{ + OptionalHeader: &astits.PESOptionalHeader{ + MarkerBits: 2, + PTSDTSIndicator: astits.PTSDTSIndicatorBothPresent, + DTS: &astits.ClockReference{Base: int64(dts.Seconds() * 90000)}, + PTS: &astits.ClockReference{Base: int64(pts.Seconds() * 90000)}, + }, + StreamID: 224, // video + }, + Data: enc, + }, + }) + if err != nil { + return err + } + + fmt.Println("wrote ts packet") + return nil +} diff --git a/examples/client-read-h264/main.go b/examples/client-read-h264/main.go index c56e5885..96baf6ee 100644 --- a/examples/client-read-h264/main.go +++ b/examples/client-read-h264/main.go @@ -26,7 +26,7 @@ func main() { return } - // parse RTP packets + // parse RTP packet var pkt rtp.Packet err := pkt.Unmarshal(payload) if err != nil { @@ -80,9 +80,8 @@ func main() { return -1 }() if h264Track < 0 { - panic(fmt.Errorf("H264 track not found")) + panic("H264 track not found") } - fmt.Printf("H264 track is number %d\n", h264Track+1) // setup all tracks for _, t := range tracks {