update H265 track parameters when resolution, VPS, SPS or PPS change

bluenviron · Dec 14, 2022 · 39da300 · 39da300
1 parent 4bafa4e
commit 39da300
Show file tree

Hide file tree

Showing 5 changed files with 262 additions and 105 deletions.
diff --git a/go.mod b/go.mod
@@ -5,7 +5,7 @@ go 1.18
 require (
 	code.cloudfoundry.org/bytefmt v0.0.0-20211005130812-5bb3c17173e5
 	github.com/abema/go-mp4 v0.8.0
-	github.com/aler9/gortsplib/v2 v2.0.0-20221213201904-04d1de717768
+	github.com/aler9/gortsplib/v2 v2.0.0-20221214135702-6141afcfc4c5
 	github.com/asticode/go-astits v1.10.1-0.20220319093903-4abe66a9b757
 	github.com/fsnotify/fsnotify v1.4.9
 	github.com/gin-gonic/gin v1.8.1

diff --git a/go.sum b/go.sum
@@ -6,8 +6,8 @@ github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 h1:JYp7IbQjafo
 github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d h1:UQZhZ2O0vMHr2cI+DC1Mbh0TJxzA3RcLoMsFw+aXw7E=
 github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho=
-github.com/aler9/gortsplib/v2 v2.0.0-20221213201904-04d1de717768 h1:Q7Qm3HjO8bYo0P2dPmKXYbfvpqhA9ZJuJ+TqN6UcdqI=
-github.com/aler9/gortsplib/v2 v2.0.0-20221213201904-04d1de717768/go.mod h1:zJ+fWtakOMN6cKV169EMNVBLPTITArrJKu/fyM+dov8=
+github.com/aler9/gortsplib/v2 v2.0.0-20221214135702-6141afcfc4c5 h1:QHF5GXc8Q23sCYYAfD24mwOJt1CX9by3KY5gYolYEaw=
+github.com/aler9/gortsplib/v2 v2.0.0-20221214135702-6141afcfc4c5/go.mod h1:zJ+fWtakOMN6cKV169EMNVBLPTITArrJKu/fyM+dov8=
 github.com/aler9/writerseeker v0.0.0-20220601075008-6f0e685b9c82 h1:9WgSzBLo3a9ToSVV7sRTBYZ1GGOZUpq4+5H3SN0UZq4=
 github.com/aler9/writerseeker v0.0.0-20220601075008-6f0e685b9c82/go.mod h1:qsMrZCbeBf/mCLOeF16KDkPu4gktn/pOWyaq1aYQE7U=
 github.com/asticode/go-astikit v0.20.0 h1:+7N+J4E4lWx2QOkRdOf6DafWJMv6O4RRfgClwQokrH8=

diff --git a/internal/core/formatprocessor_h264.go b/internal/core/formatprocessor_h264.go
@@ -10,24 +10,9 @@ import (
 	"github.com/pion/rtp"
 )
 
-type dataH264 struct {
-	rtpPackets []*rtp.Packet
-	ntp        time.Time
-	pts        time.Duration
-	nalus      [][]byte
-}
-
-func (d *dataH264) getRTPPackets() []*rtp.Packet {
-	return d.rtpPackets
-}
-
-func (d *dataH264) getNTP() time.Time {
-	return d.ntp
-}
-
 // extract SPS and PPS without decoding RTP packets
 func rtpH264ExtractSPSPPS(pkt *rtp.Packet) ([]byte, []byte) {
-	if len(pkt.Payload) == 0 {
+	if len(pkt.Payload) < 1 {
 		return nil, nil
 	}
 
@@ -40,7 +25,7 @@ func rtpH264ExtractSPSPPS(pkt *rtp.Packet) ([]byte, []byte) {
 	case h264.NALUTypePPS:
 		return nil, pkt.Payload
 
-	case 24: // STAP-A
+	case h264.NALUTypeSTAPA:
 		payload := pkt.Payload[1:]
 		var sps []byte
 		var pps []byte
@@ -53,10 +38,14 @@ func rtpH264ExtractSPSPPS(pkt *rtp.Packet) ([]byte, []byte) {
 			size := uint16(payload[0])<<8 | uint16(payload[1])
 			payload = payload[2:]
 
-			if size == 0 || int(size) > len(payload) {
+			if size == 0 {
 				break
 			}
 
+			if int(size) > len(payload) {
+				return nil, nil
+			}
+
 			nalu := payload[:size]
 			payload = payload[size:]
 
@@ -78,6 +67,21 @@ func rtpH264ExtractSPSPPS(pkt *rtp.Packet) ([]byte, []byte) {
 	}
 }
 
+type dataH264 struct {
+	rtpPackets []*rtp.Packet
+	ntp        time.Time
+	pts        time.Duration
+	nalus      [][]byte
+}
+
+func (d *dataH264) getRTPPackets() []*rtp.Packet {
+	return d.rtpPackets
+}
+
+func (d *dataH264) getNTP() time.Time {
+	return d.ntp
+}
+
 type formatProcessorH264 struct {
 	format *format.H264
 

diff --git a/internal/core/formatprocessor_h265.go b/internal/core/formatprocessor_h265.go
@@ -1,14 +1,80 @@
 package core
 
 import (
+	"bytes"
 	"fmt"
 	"time"
 
 	"github.com/aler9/gortsplib/v2/pkg/format"
 	"github.com/aler9/gortsplib/v2/pkg/formatdecenc/rtph265"
+	"github.com/aler9/gortsplib/v2/pkg/h265"
 	"github.com/pion/rtp"
 )
 
+// extract VPS, SPS and PPS without decoding RTP packets
+func rtpH265ExtractVPSSPSPPS(pkt *rtp.Packet) ([]byte, []byte, []byte) {
+	if len(pkt.Payload) < 2 {
+		return nil, nil, nil
+	}
+
+	typ := h265.NALUType((pkt.Payload[0] >> 1) & 0b111111)
+
+	switch typ {
+	case h265.NALUTypeVPS:
+		return pkt.Payload, nil, nil
+
+	case h265.NALUTypeSPS:
+		return nil, pkt.Payload, nil
+
+	case h265.NALUTypePPS:
+		return nil, nil, pkt.Payload
+
+	case h265.NALUTypeAggregationUnit:
+		payload := pkt.Payload[2:]
+		var vps []byte
+		var sps []byte
+		var pps []byte
+
+		for len(payload) > 0 {
+			if len(payload) < 2 {
+				break
+			}
+
+			size := uint16(payload[0])<<8 | uint16(payload[1])
+			payload = payload[2:]
+
+			if size == 0 {
+				break
+			}
+
+			if int(size) > len(payload) {
+				return nil, nil, nil
+			}
+
+			nalu := payload[:size]
+			payload = payload[size:]
+
+			typ = h265.NALUType((pkt.Payload[0] >> 1) & 0b111111)
+
+			switch typ {
+			case h265.NALUTypeVPS:
+				vps = nalu
+
+			case h265.NALUTypeSPS:
+				sps = nalu
+
+			case h265.NALUTypePPS:
+				pps = nalu
+			}
+		}
+
+		return vps, sps, pps
+
+	default:
+		return nil, nil, nil
+	}
+}
+
 type dataH265 struct {
 	rtpPackets []*rtp.Packet
 	ntp        time.Time
@@ -47,7 +113,19 @@ func newFormatProcessorH265(
 }
 
 func (t *formatProcessorH265) updateTrackParametersFromRTPPacket(pkt *rtp.Packet) {
-	// TODO: extract VPS, SPS, PPS and set them into the track
+	vps, sps, pps := rtpH265ExtractVPSSPSPPS(pkt)
+
+	if vps != nil && !bytes.Equal(vps, t.format.SafeVPS()) {
+		t.format.SafeSetVPS(vps)
+	}
+
+	if sps != nil && !bytes.Equal(sps, t.format.SafeSPS()) {
+		t.format.SafeSetSPS(sps)
+	}
+
+	if pps != nil && !bytes.Equal(pps, t.format.SafePPS()) {
+		t.format.SafeSetPPS(pps)
+	}
 }
 
 func (t *formatProcessorH265) updateTrackParametersFromNALUs(nalus [][]byte) {