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
-
-	// 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 {