h265: fix DTS extraction of streams with short_term_ref_pic_set_sps_f…

…lag=1 (bluenviron/mediamtx#2417)
bluenviron · Oct 10, 2023 · dfc514b · dfc514b
1 parent 4c7d799
commit dfc514b
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 7 deletions.
diff --git a/pkg/codecs/h265/dts_extractor.go b/pkg/codecs/h265/dts_extractor.go
@@ -3,6 +3,7 @@ package h265
 import (
 	"fmt"
 	"time"
+	"math"
 
 	"github.com/bluenviron/mediacommon/pkg/bits"
 	"github.com/bluenviron/mediacommon/pkg/codecs/h264"
@@ -78,14 +79,31 @@ func getPTSDTSDiff(buf []byte, sps *SPS, pps *PPS) (uint32, error) {
 		return 0, err
 	}
 
-	if shortTermRefPicSetSpsFlag {
-		return 0, fmt.Errorf("short_term_ref_pic_set_sps_flag = true is not supported")
-	}
+	var rps *SPS_ShortTermRefPicSet
 
-	var rps SPS_ShortTermRefPicSet
-	err = rps.unmarshal(buf, &pos, uint32(len(sps.ShortTermRefPicSets)), uint32(len(sps.ShortTermRefPicSets)), nil)
-	if err != nil {
-		return 0, err
+	if !shortTermRefPicSetSpsFlag {
+		rps = &SPS_ShortTermRefPicSet{}
+		err = rps.unmarshal(buf, &pos, uint32(len(sps.ShortTermRefPicSets)), uint32(len(sps.ShortTermRefPicSets)), nil)
+		if err != nil {
+			return 0, err
+		}
+	} else {
+		if len(sps.ShortTermRefPicSets) == 0 {
+			return 0, fmt.Errorf("invalid short_term_ref_pic_set_idx")
+		}
+
+		b := int(math.Ceil(math.Log2(float64(len(sps.ShortTermRefPicSets)))))
+		tmp, err := bits.ReadBits(buf, &pos, b)
+		if err != nil {
+			return 0, err
+		}
+		shortTermRefPicSetIdx := int(tmp)
+
+		if len(sps.ShortTermRefPicSets) <= shortTermRefPicSetIdx {
+			return 0, fmt.Errorf("invalid short_term_ref_pic_set_idx")
+		}
+
+		rps = sps.ShortTermRefPicSets[shortTermRefPicSetIdx]
 	}
 
 	var v uint32

diff --git a/pkg/codecs/h265/dts_extractor_test.go b/pkg/codecs/h265/dts_extractor_test.go
@@ -129,6 +129,61 @@ func TestDTSExtractor(t *testing.T) {
 				},
 			},
 		},
+		{
+			"short_term_ref_pic_set_sps_flag",
+			[]sequenceSample{
+				{
+					[][]byte{
+						{ // VPS
+							0x40, 0x01, 0x0c, 0x01, 0xff, 0xff, 0x01, 0x40,
+							0x00, 0x00, 0x03, 0x00, 0x80, 0x00, 0x00, 0x03,
+							0x00, 0x00, 0x03, 0x00, 0x99, 0xa5, 0x02, 0x40,
+						},
+						{ // SPS
+							0x42, 0x01, 0x01, 0x01, 0x40, 0x00, 0x00, 0x03,
+							0x00, 0x80, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03,
+							0x00, 0x99, 0xa0, 0x03, 0xc0, 0x80, 0x10, 0xe5,
+							0x8d, 0xa5, 0x92, 0x42, 0x36, 0x22, 0xec, 0xb8,
+							0x80, 0x40, 0x00, 0x00, 0x03, 0x00, 0x40, 0x00,
+							0x00, 0x05, 0x0f, 0xe2, 0xc4, 0xa0,
+						},
+						{ // PPS
+							0x44, 0x01, 0xc0, 0xe0, 0x98, 0x93, 0x03, 0x05,
+							0x14, 0x90,
+						},
+						{ // IDR
+							0x26, 0x01, 0xaf, 0x3e, 0x3d, 0x3a, 0xca, 0xc0,
+							0xf2, 0x2f, 0xc3, 0x0f, 0x86, 0x9f, 0xed, 0xfc,
+							0x67, 0x2f, 0x62, 0x69,
+						},
+					},
+					1113436 * time.Nanosecond,
+					-48886564 * time.Nanosecond,
+				},
+				{
+					[][]byte{
+						{ // TRAIL_R
+							0x02, 0x02, 0xd0, 0x00, 0x0c, 0xc6, 0x27, 0xfe,
+							0x6e, 0x6d, 0xe8, 0x10, 0xd5, 0xce, 0x61, 0x1b,
+							0x66, 0xf6, 0x21, 0x59,
+						},
+					},
+					68113436 * time.Nanosecond,
+					18113436 * time.Nanosecond,
+				},
+				{
+					[][]byte{
+						{ // TRAIL_R
+							0x02, 0x02, 0xd0, 0x00, 0x14, 0xc6, 0x7c, 0xfe,
+							0x83, 0x29, 0x34, 0xba, 0xce, 0xaa, 0x8b, 0x76,
+							0xb0, 0x95, 0x67, 0xb2,
+						},
+					},
+					101113436 * time.Nanosecond,
+					51113436 * time.Nanosecond,
+				},
+			},
+		},
 	} {
 		t.Run(ca.name, func(t *testing.T) {
 			ex := NewDTSExtractor()