cleanup examples

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
+	var enc *mpegtsEncoder
-
-	// add an H264 track to the MPEG-TS muxer
-	mux.AddElementaryStream(astits.PMTElementaryStream{
-		ElementaryPID: 256,
-		StreamType:    astits.StreamTypeH264Video,
-	})
-	mux.SetPCRPID(256)
 
 	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 {
+			// encode H264 NALUs into MPEG-TS
-				firstPacketWritten = true
+			err = enc.encode(nalus, pts)
-				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)
 			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
+	// setup the encoder
-	dec = rtph264.NewDecoder()
+	enc, err = newMPEGTSEncoder(h264Conf)
+	if err != nil {
+		panic(err)
+	}
 
 	// setup all tracks
 	for _, t := range tracks {

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
+}

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 {