add client-read-save-to-disk example

README.md

@@ -50,6 +50,7 @@ Features:
 * [client-read-options](examples/client-read-options/main.go)
 * [client-read-pause](examples/client-read-pause/main.go)
 * [client-read-h264](examples/client-read-h264/main.go)
+* [client-read-save-to.disk](examples/client-read-save-to-disk/main.go)
 * [client-publish](examples/client-publish/main.go)
 * [client-publish-options](examples/client-publish-options/main.go)
 * [client-publish-pause](examples/client-publish-pause/main.go)

examples/client-read-save-to-disk/main.go

@@ -0,0 +1,184 @@
+package main
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/aler9/gortsplib"
+	"github.com/aler9/gortsplib/pkg/h264"
+	"github.com/aler9/gortsplib/pkg/rtph264"
+	"github.com/asticode/go-astits"
+	"github.com/pion/rtp"
+)
+
+// This example shows how to
+// 1. connect to a RTSP server and read all tracks on a path
+// 2. check whether there's a H264 track
+// 3. save the content of the H264 track to a file in MPEG-TS format
+
+const (
+	inputStream = "rtsp://localhost:8554/mystream"
+	outputFile  = "mystream.ts"
+)
+
+func main() {
+	p := gortsplib.ClientProtocolUDP
+	c := gortsplib.Client{
+		Protocol: &p,
+	}
+
+	// connect to the server and start reading all tracks
+	conn, err := c.DialRead(inputStream)
+	if err != nil {
+		panic(err)
+	}
+	defer conn.Close()
+
+	// find the H264 track
+	var h264TrackID int = -1
+	var h264Conf *gortsplib.TrackConfigH264
+	for i, track := range conn.Tracks() {
+		if track.IsH264() {
+			h264TrackID = i
+			h264Conf, err = track.ExtractConfigH264()
+			if err != nil {
+				panic(err)
+			}
+		}
+	}
+	if h264TrackID < 0 {
+		panic(fmt.Errorf("H264 track not found"))
+	}
+
+	// open output file
+	f, err := os.Create(outputFile)
+	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 naluBuffer [][]byte
+	var startPTS time.Duration
+
+	// add an H264 track to the MPEG-TS muxer
+	mux.AddElementaryStream(astits.PMTElementaryStream{
+		ElementaryPID: 256,
+		StreamType:    astits.StreamTypeH264Video,
+	})
+	mux.SetPCRPID(256)
+
+	// read RTP frames
+	err = conn.ReadFrames(func(trackID int, streamType gortsplib.StreamType, payload []byte) {
+		if trackID != h264TrackID {
+			return
+		}
+
+		if streamType != gortsplib.StreamTypeRTP {
+			return
+		}
+
+		// parse RTP packets
+		var pkt rtp.Packet
+		err := pkt.Unmarshal(payload)
+		if err != nil {
+			return
+		}
+
+		// convert RTP frames into H264 NALUs
+		nalus, pts, err := dec.DecodeRTP(&pkt)
+		if err != nil {
+			return
+		}
+
+		// store NALUs in a buffer until a packet with the Marker flag is received
+		naluBuffer = append(naluBuffer, nalus...)
+
+		// RTP marker means that all the NALUs with the same PTS have been received.
+		if pkt.Marker {
+			if !firstPacketWritten {
+				firstPacketWritten = true
+				startPTS = pts
+			}
+
+			// check whether there's an IDR
+			idrPresent := func() bool {
+				for _, nalu := range naluBuffer {
+					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 naluBuffer {
+				// 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)
+			}
+
+			naluBuffer = nil
+
+			// encode into Annex-B
+			enc, err := h264.EncodeAnnexB(filteredNALUs)
+			if err != nil {
+				panic(err)
+			}
+
+			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")
+		}
+	})
+	panic(err)
+}