From 3ee5fbdcf753aa4683c8419a0c310b580d628ee5 Mon Sep 17 00:00:00 2001 From: Alessandro Ros Date: Fri, 21 Jun 2024 23:01:46 +0200 Subject: [PATCH] add client-play-format-h264-mpeg4audio-save-to-disk example (#230) (#582) --- README.md | 1 + .../main.go | 146 ++++++++++++++++++ .../mpegts_muxer.go | 134 ++++++++++++++++ .../main.go | 5 +- .../mpegts_muxer.go | 2 +- .../main.go | 5 +- .../mpegts_muxer.go | 2 +- .../main.go | 5 +- .../main.go | 12 +- examples/server-h264-save-to-disk/main.go | 4 +- .../server-h264-save-to-disk/mpegts_muxer.go | 2 +- 11 files changed, 299 insertions(+), 19 deletions(-) create mode 100644 examples/client-play-format-h264-mpeg4audio-save-to-disk/main.go create mode 100644 examples/client-play-format-h264-mpeg4audio-save-to-disk/mpegts_muxer.go diff --git a/README.md b/README.md index ff7052c0..1deb2ec0 100644 --- a/README.md +++ b/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) diff --git a/examples/client-play-format-h264-mpeg4audio-save-to-disk/main.go b/examples/client-play-format-h264-mpeg4audio-save-to-disk/main.go new file mode 100644 index 00000000..d0283708 --- /dev/null +++ b/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()) +} diff --git a/examples/client-play-format-h264-mpeg4audio-save-to-disk/mpegts_muxer.go b/examples/client-play-format-h264-mpeg4audio-save-to-disk/mpegts_muxer.go new file mode 100644 index 00000000..b5b51699 --- /dev/null +++ b/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) +} diff --git a/examples/client-play-format-h264-save-to-disk/main.go b/examples/client-play-format-h264-save-to-disk/main.go index 8f92fa39..e545b0da 100644 --- a/examples/client-play-format-h264-save-to-disk/main.go +++ b/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) diff --git a/examples/client-play-format-h264-save-to-disk/mpegts_muxer.go b/examples/client-play-format-h264-save-to-disk/mpegts_muxer.go index 4339bc75..62d03fb7 100644 --- a/examples/client-play-format-h264-save-to-disk/mpegts_muxer.go +++ b/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) } diff --git a/examples/client-play-format-h265-save-to-disk/main.go b/examples/client-play-format-h265-save-to-disk/main.go index 0481519a..e737ccb2 100644 --- a/examples/client-play-format-h265-save-to-disk/main.go +++ b/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) diff --git a/examples/client-play-format-h265-save-to-disk/mpegts_muxer.go b/examples/client-play-format-h265-save-to-disk/mpegts_muxer.go index 71ed13c9..fd94ae77 100644 --- a/examples/client-play-format-h265-save-to-disk/mpegts_muxer.go +++ b/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) } diff --git a/examples/client-play-format-mpeg4audio-save-to-disk/main.go b/examples/client-play-format-mpeg4audio-save-to-disk/main.go index 31f0e08e..753e91c4 100644 --- a/examples/client-play-format-mpeg4audio-save-to-disk/main.go +++ b/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) diff --git a/examples/client-play-format-opus-save-to-disk/main.go b/examples/client-play-format-opus-save-to-disk/main.go index 4e2184f7..925fb61f 100644 --- a/examples/client-play-format-opus-save-to-disk/main.go +++ b/examples/client-play-format-opus-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 Opus 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{} @@ -55,19 +55,15 @@ func main() { fileName: "mystream.ts", track: &mpegts.Track{ Codec: &mpegts.CodecOpus{ - ChannelCount: func() int { - if forma.IsStereo { - return 2 - } - return 1 - }(), + ChannelCount: forma.ChannelCount, }, }, } - 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) diff --git a/examples/server-h264-save-to-disk/main.go b/examples/server-h264-save-to-disk/main.go index f4ba5bec..d29e9da6 100644 --- a/examples/server-h264-save-to-disk/main.go +++ b/examples/server-h264-save-to-disk/main.go @@ -17,7 +17,7 @@ import ( // This example shows how to // 1. create a RTSP server which accepts plain connections // 2. allow a single client to publish a stream, containing a H264 media, with TCP or UDP -// 3. save the content of the H264 media into a file in MPEG-TS format +// 3. save the content of the H264 media in a file in MPEG-TS format type serverHandler struct { s *gortsplib.Server @@ -88,7 +88,7 @@ func (sh *serverHandler) OnAnnounce(ctx *gortsplib.ServerHandlerOnAnnounceCtx) ( sps: forma.SPS, pps: forma.PPS, } - mpegtsMuxer.initialize() + err = mpegtsMuxer.initialize() if err != nil { return &base.Response{ StatusCode: base.StatusBadRequest, diff --git a/examples/server-h264-save-to-disk/mpegts_muxer.go b/examples/server-h264-save-to-disk/mpegts_muxer.go index 9b198982..dcfe268b 100644 --- a/examples/server-h264-save-to-disk/mpegts_muxer.go +++ b/examples/server-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.WriteH265(e.track, durationGoToMPEGTS(pts), durationGoToMPEGTS(dts), idrPresent, au) }