From d40f9a53c959c40c92a63901cc1619cf661d02d3 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Sat, 29 Oct 2022 17:07:09 +0200
Subject: [PATCH] examples: insert SPS/PPS before IDR groups once

---
 .../mpegtsmuxer.go                             | 18 ++++++++++--------
 .../server-h264-save-to-disk/mpegtsmuxer.go    | 18 ++++++++++--------
 2 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/examples/client-read-h264-save-to-disk/mpegtsmuxer.go b/examples/client-read-h264-save-to-disk/mpegtsmuxer.go
index 5f2f01a5..53b3d31c 100644
--- a/examples/client-read-h264-save-to-disk/mpegtsmuxer.go
+++ b/examples/client-read-h264-save-to-disk/mpegtsmuxer.go
@@ -81,11 +81,6 @@ func (e *mpegtsMuxer) encode(nalus [][]byte, pts time.Duration) error {
 		case h264.NALUTypeIDR:
 			idrPresent = true
 
-			// add SPS and PPS before every IDR
-			if e.sps != nil && e.pps != nil {
-				filteredNALUs = append(filteredNALUs, e.sps, e.pps)
-			}
-
 		case h264.NALUTypeNonIDR:
 			nonIDRPresent = true
 		}
@@ -93,10 +88,17 @@ func (e *mpegtsMuxer) encode(nalus [][]byte, pts time.Duration) error {
 		filteredNALUs = append(filteredNALUs, nalu)
 	}
 
+	nalus = filteredNALUs
+
 	if !nonIDRPresent && !idrPresent {
 		return nil
 	}
 
+	// add SPS and PPS before every group that contains an IDR
+	if idrPresent {
+		nalus = append([][]byte{e.sps, e.pps}, nalus...)
+	}
+
 	var dts time.Duration
 
 	if !e.firstIDRReceived {
@@ -109,7 +111,7 @@ func (e *mpegtsMuxer) encode(nalus [][]byte, pts time.Duration) error {
 		e.dtsExtractor = h264.NewDTSExtractor()
 
 		var err error
-		dts, err = e.dtsExtractor.Extract(filteredNALUs, pts)
+		dts, err = e.dtsExtractor.Extract(nalus, pts)
 		if err != nil {
 			return err
 		}
@@ -120,7 +122,7 @@ func (e *mpegtsMuxer) encode(nalus [][]byte, pts time.Duration) error {
 
 	} else {
 		var err error
-		dts, err = e.dtsExtractor.Extract(filteredNALUs, pts)
+		dts, err = e.dtsExtractor.Extract(nalus, pts)
 		if err != nil {
 			return err
 		}
@@ -143,7 +145,7 @@ func (e *mpegtsMuxer) encode(nalus [][]byte, pts time.Duration) error {
 	}
 
 	// encode into Annex-B
-	annexb, err := h264.AnnexBMarshal(filteredNALUs)
+	annexb, err := h264.AnnexBMarshal(nalus)
 	if err != nil {
 		return err
 	}
diff --git a/examples/server-h264-save-to-disk/mpegtsmuxer.go b/examples/server-h264-save-to-disk/mpegtsmuxer.go
index 5f2f01a5..53b3d31c 100644
--- a/examples/server-h264-save-to-disk/mpegtsmuxer.go
+++ b/examples/server-h264-save-to-disk/mpegtsmuxer.go
@@ -81,11 +81,6 @@ func (e *mpegtsMuxer) encode(nalus [][]byte, pts time.Duration) error {
 		case h264.NALUTypeIDR:
 			idrPresent = true
 
-			// add SPS and PPS before every IDR
-			if e.sps != nil && e.pps != nil {
-				filteredNALUs = append(filteredNALUs, e.sps, e.pps)
-			}
-
 		case h264.NALUTypeNonIDR:
 			nonIDRPresent = true
 		}
@@ -93,10 +88,17 @@ func (e *mpegtsMuxer) encode(nalus [][]byte, pts time.Duration) error {
 		filteredNALUs = append(filteredNALUs, nalu)
 	}
 
+	nalus = filteredNALUs
+
 	if !nonIDRPresent && !idrPresent {
 		return nil
 	}
 
+	// add SPS and PPS before every group that contains an IDR
+	if idrPresent {
+		nalus = append([][]byte{e.sps, e.pps}, nalus...)
+	}
+
 	var dts time.Duration
 
 	if !e.firstIDRReceived {
@@ -109,7 +111,7 @@ func (e *mpegtsMuxer) encode(nalus [][]byte, pts time.Duration) error {
 		e.dtsExtractor = h264.NewDTSExtractor()
 
 		var err error
-		dts, err = e.dtsExtractor.Extract(filteredNALUs, pts)
+		dts, err = e.dtsExtractor.Extract(nalus, pts)
 		if err != nil {
 			return err
 		}
@@ -120,7 +122,7 @@ func (e *mpegtsMuxer) encode(nalus [][]byte, pts time.Duration) error {
 
 	} else {
 		var err error
-		dts, err = e.dtsExtractor.Extract(filteredNALUs, pts)
+		dts, err = e.dtsExtractor.Extract(nalus, pts)
 		if err != nil {
 			return err
 		}
@@ -143,7 +145,7 @@ func (e *mpegtsMuxer) encode(nalus [][]byte, pts time.Duration) error {
 	}
 
 	// encode into Annex-B
-	annexb, err := h264.AnnexBMarshal(filteredNALUs)
+	annexb, err := h264.AnnexBMarshal(nalus)
 	if err != nil {
 		return err
 	}