add client-play-format-h264-mpeg4audio-save-to-disk example (#230) (#582

)

README.md

@@ -66,6 +66,7 @@ Features:
 * [client-play-format-h264](examples/client-play-format-h264/main.go)
 * [client-play-format-h264-convert-to-jpeg](examples/client-play-format-h264-convert-to-jpeg/main.go)
 * [client-play-format-h264-save-to-disk](examples/client-play-format-h264-save-to-disk/main.go)
+* [client-play-format-h264-mpeg4audio-save-to-disk](examples/client-play-format-h264-mpeg4audio-save-to-disk/main.go)
 * [client-play-format-h265](examples/client-play-format-h265/main.go)
 * [client-play-format-h265-convert-to-jpeg](examples/client-play-format-h265-convert-to-jpeg/main.go)
 * [client-play-format-h265-save-to-disk](examples/client-play-format-h265-save-to-disk/main.go)

examples/client-play-format-h264-mpeg4audio-save-to-disk/main.go

@@ -0,0 +1,146 @@
+package main
+
+import (
+	"log"
+
+	"github.com/bluenviron/gortsplib/v4"
+	"github.com/bluenviron/gortsplib/v4/pkg/base"
+	"github.com/bluenviron/gortsplib/v4/pkg/format"
+	"github.com/bluenviron/gortsplib/v4/pkg/format/rtph264"
+	"github.com/pion/rtp"
+)
+
+// This example shows how to
+// 1. connect to a RTSP server
+// 2. check if there's a H264 format and an MPEG-4 audio format
+// 3. save the content of those formats in a file in MPEG-TS format
+
+func main() {
+	c := gortsplib.Client{}
+
+	// parse URL
+	u, err := base.ParseURL("rtsp://localhost:8554/stream")
+	if err != nil {
+		panic(err)
+	}
+
+	// connect to the server
+	err = c.Start(u.Scheme, u.Host)
+	if err != nil {
+		panic(err)
+	}
+	defer c.Close()
+
+	// find available medias
+	desc, _, err := c.Describe(u)
+	if err != nil {
+		panic(err)
+	}
+
+	// find the H264 media and format
+	var h264Format *format.H264
+	h264Media := desc.FindFormat(&h264Format)
+	if h264Media == nil {
+		panic("H264 media not found")
+	}
+
+	// find the MPEG-4 audio media and format
+	var mpeg4AudioFormat *format.MPEG4Audio
+	mpeg4AudioMedia := desc.FindFormat(&mpeg4AudioFormat)
+	if mpeg4AudioMedia == nil {
+		panic("MPEG-4 audio media not found")
+	}
+
+	// setup RTP -> H264 decoder
+	h264RTPDec, err := h264Format.CreateDecoder()
+	if err != nil {
+		panic(err)
+	}
+
+	// setup RTP -> MPEG-4 audio decoder
+	mpeg4AudioRTPDec, err := mpeg4AudioFormat.CreateDecoder()
+	if err != nil {
+		panic(err)
+	}
+
+	// setup MPEG-TS muxer
+	mpegtsMuxer := &mpegtsMuxer{
+		fileName:         "mystream.ts",
+		h264Format:       h264Format,
+		mpeg4AudioFormat: mpeg4AudioFormat,
+	}
+	err = mpegtsMuxer.initialize()
+	if err != nil {
+		panic(err)
+	}
+	defer mpegtsMuxer.close()
+
+	// setup all medias
+	err = c.SetupAll(desc.BaseURL, desc.Medias)
+	if err != nil {
+		panic(err)
+	}
+
+	// called when a H264/RTP packet arrives
+	c.OnPacketRTP(h264Media, h264Format, func(pkt *rtp.Packet) {
+		// decode timestamp
+		pts, ok := c.PacketPTS(h264Media, pkt)
+		if !ok {
+			log.Printf("waiting for timestamp")
+			return
+		}
+
+		// extract access unit from RTP packets
+		au, err := h264RTPDec.Decode(pkt)
+		if err != nil {
+			if err != rtph264.ErrNonStartingPacketAndNoPrevious && err != rtph264.ErrMorePacketsNeeded {
+				log.Printf("ERR: %v", err)
+			}
+			return
+		}
+
+		// encode the access unit into MPEG-TS
+		err = mpegtsMuxer.writeH264(au, pts)
+		if err != nil {
+			log.Printf("ERR: %v", err)
+			return
+		}
+
+		log.Printf("saved TS packet")
+	})
+
+	// called when a MPEG-4 audio / RTP packet arrives
+	c.OnPacketRTP(mpeg4AudioMedia, mpeg4AudioFormat, func(pkt *rtp.Packet) {
+		// decode timestamp
+		pts, ok := c.PacketPTS(mpeg4AudioMedia, pkt)
+		if !ok {
+			log.Printf("waiting for timestamp")
+			return
+		}
+
+		// extract access units from RTP packets
+		aus, err := mpeg4AudioRTPDec.Decode(pkt)
+		if err != nil {
+			log.Printf("ERR: %v", err)
+			return
+		}
+
+		// encode access units into MPEG-TS
+		err = mpegtsMuxer.writeMPEG4Audio(aus, pts)
+		if err != nil {
+			log.Printf("ERR: %v", err)
+			return
+		}
+
+		log.Printf("saved TS packet")
+	})
+
+	// start playing
+	_, err = c.Play(nil)
+	if err != nil {
+		panic(err)
+	}
+
+	// wait until a fatal error
+	panic(c.Wait())
+}

examples/client-play-format-h264-mpeg4audio-save-to-disk/mpegts_muxer.go

@@ -0,0 +1,134 @@
+package main
+
+import (
+	"bufio"
+	"os"
+	"sync"
+	"time"
+
+	"github.com/bluenviron/gortsplib/v4/pkg/format"
+	"github.com/bluenviron/mediacommon/pkg/codecs/h264"
+	"github.com/bluenviron/mediacommon/pkg/formats/mpegts"
+)
+
+func durationGoToMPEGTS(v time.Duration) int64 {
+	return int64(v.Seconds() * 90000)
+}
+
+// mpegtsMuxer allows to save a H264 / MPEG-4 audio stream into a MPEG-TS file.
+type mpegtsMuxer struct {
+	fileName         string
+	h264Format       *format.H264
+	mpeg4AudioFormat *format.MPEG4Audio
+
+	f               *os.File
+	b               *bufio.Writer
+	w               *mpegts.Writer
+	h264Track       *mpegts.Track
+	mpeg4AudioTrack *mpegts.Track
+	dtsExtractor    *h264.DTSExtractor
+	mutex           sync.Mutex
+}
+
+// initialize initializes a mpegtsMuxer.
+func (e *mpegtsMuxer) initialize() error {
+	var err error
+	e.f, err = os.Create(e.fileName)
+	if err != nil {
+		return err
+	}
+	e.b = bufio.NewWriter(e.f)
+
+	e.h264Track = &mpegts.Track{
+		Codec: &mpegts.CodecH264{},
+	}
+
+	e.mpeg4AudioTrack = &mpegts.Track{
+		Codec: &mpegts.CodecMPEG4Audio{
+			Config: *e.mpeg4AudioFormat.Config,
+		},
+	}
+
+	e.w = mpegts.NewWriter(e.b, []*mpegts.Track{e.h264Track, e.mpeg4AudioTrack})
+
+	return nil
+}
+
+// close closes all the mpegtsMuxer resources.
+func (e *mpegtsMuxer) close() {
+	e.b.Flush()
+	e.f.Close()
+}
+
+// writeH264 writes a H264 access unit into MPEG-TS.
+func (e *mpegtsMuxer) writeH264(au [][]byte, pts time.Duration) error {
+	e.mutex.Lock()
+	defer e.mutex.Unlock()
+
+	var filteredAU [][]byte
+
+	nonIDRPresent := false
+	idrPresent := false
+
+	for _, nalu := range au {
+		typ := h264.NALUType(nalu[0] & 0x1F)
+		switch typ {
+		case h264.NALUTypeSPS:
+			e.h264Format.SPS = nalu
+			continue
+
+		case h264.NALUTypePPS:
+			e.h264Format.PPS = nalu
+			continue
+
+		case h264.NALUTypeAccessUnitDelimiter:
+			continue
+
+		case h264.NALUTypeIDR:
+			idrPresent = true
+
+		case h264.NALUTypeNonIDR:
+			nonIDRPresent = true
+		}
+
+		filteredAU = append(filteredAU, nalu)
+	}
+
+	au = filteredAU
+
+	if au == nil || (!nonIDRPresent && !idrPresent) {
+		return nil
+	}
+
+	// add SPS and PPS before access unit that contains an IDR
+	if idrPresent {
+		au = append([][]byte{e.h264Format.SPS, e.h264Format.PPS}, au...)
+	}
+
+	var dts time.Duration
+
+	if e.dtsExtractor == nil {
+		// skip samples silently until we find one with a IDR
+		if !idrPresent {
+			return nil
+		}
+		e.dtsExtractor = h264.NewDTSExtractor()
+	}
+
+	var err error
+	dts, err = e.dtsExtractor.Extract(au, pts)
+	if err != nil {
+		return err
+	}
+
+	// encode into MPEG-TS
+	return e.w.WriteH264(e.h264Track, durationGoToMPEGTS(pts), durationGoToMPEGTS(dts), idrPresent, au)
+}
+
+// writeMPEG4Audio writes MPEG-4 audio access units into MPEG-TS.
+func (e *mpegtsMuxer) writeMPEG4Audio(aus [][]byte, pts time.Duration) error {
+	e.mutex.Lock()
+	defer e.mutex.Unlock()
+
+	return e.w.WriteMPEG4Audio(e.mpeg4AudioTrack, durationGoToMPEGTS(pts), aus)
+}

examples/client-play-format-h264-save-to-disk/main.go

@@ -13,7 +13,7 @@ import (
 // This example shows how to
 // 1. connect to a RTSP server
 // 2. check if there's a H264 format
-// 3. save the content of the format into a file in MPEG-TS format
+// 3. save the content of the format in a file in MPEG-TS format
 
 func main() {
 	c := gortsplib.Client{}
@@ -56,10 +56,11 @@ func main() {
 		sps:      forma.SPS,
 		pps:      forma.PPS,
 	}
-	mpegtsMuxer.initialize()
+	err = mpegtsMuxer.initialize()
 	if err != nil {
 		panic(err)
 	}
+	defer mpegtsMuxer.close()
 
 	// setup a single media
 	_, err = c.Setup(desc.BaseURL, medi, 0, 0)

examples/client-play-format-h264-save-to-disk/mpegts_muxer.go

@@ -109,5 +109,5 @@ func (e *mpegtsMuxer) writeH264(au [][]byte, pts time.Duration) error {
 	}
 
 	// encode into MPEG-TS
-	return e.w.WriteH26x(e.track, durationGoToMPEGTS(pts), durationGoToMPEGTS(dts), idrPresent, au)
+	return e.w.WriteH264(e.track, durationGoToMPEGTS(pts), durationGoToMPEGTS(dts), idrPresent, au)
 }

examples/client-play-format-h265-save-to-disk/main.go

@@ -13,7 +13,7 @@ import (
 // This example shows how to
 // 1. connect to a RTSP server
 // 2. check if there's a H265 format
-// 3. save the content of the format into a file in MPEG-TS format
+// 3. save the content of the format in a file in MPEG-TS format
 
 func main() {
 	c := gortsplib.Client{}
@@ -57,10 +57,11 @@ func main() {
 		sps:      forma.SPS,
 		pps:      forma.PPS,
 	}
-	mpegtsMuxer.initialize()
+	err = mpegtsMuxer.initialize()
 	if err != nil {
 		panic(err)
 	}
+	defer mpegtsMuxer.close()
 
 	// setup a single media
 	_, err = c.Setup(desc.BaseURL, medi, 0, 0)

examples/client-play-format-h265-save-to-disk/mpegts_muxer.go

@@ -110,5 +110,5 @@ func (e *mpegtsMuxer) writeH265(au [][]byte, pts time.Duration) error {
 	}
 
 	// encode into MPEG-TS
-	return e.w.WriteH26x(e.track, durationGoToMPEGTS(pts), durationGoToMPEGTS(dts), isRandomAccess, au)
+	return e.w.WriteH265(e.track, durationGoToMPEGTS(pts), durationGoToMPEGTS(dts), isRandomAccess, au)
 }

examples/client-play-format-mpeg4audio-save-to-disk/main.go

@@ -13,7 +13,7 @@ import (
 // This example shows how to
 // 1. connect to a RTSP server
 // 2. check if there's an MPEG-4 audio format
-// 3. save the content of the format into a file in MPEG-TS format
+// 3. save the content of the format in a file in MPEG-TS format
 
 func main() {
 	c := gortsplib.Client{}
@@ -59,10 +59,11 @@ func main() {
 			},
 		},
 	}
-	mpegtsMuxer.initialize()
+	err = mpegtsMuxer.initialize()
 	if err != nil {
 		panic(err)
 	}
+	defer mpegtsMuxer.close()
 
 	// setup a single media
 	_, err = c.Setup(desc.BaseURL, medi, 0, 0)