codecs/av1: support AMD hardware encoder (bluenviron/mediamtx#3902) (#…

…162)
bluenviron · Dec 21, 2024 · be39cf7 · be39cf7
1 parent 7200798
commit be39cf7
Show file tree

Hide file tree

Showing 8 changed files with 120 additions and 13 deletions.
diff --git a/pkg/bits/read.go b/pkg/bits/read.go
@@ -53,7 +53,7 @@ func ReadBitsUnsafe(buf []byte, pos *int, n int) uint64 {
 // ReadGolombUnsigned reads an unsigned golomb-encoded value.
 func ReadGolombUnsigned(buf []byte, pos *int) (uint32, error) {
 	buflen := len(buf)
-	leadingZeroBits := uint32(0)
+	leadingZeros := uint32(0)
 
 	for {
 		if (buflen*8 - *pos) == 0 {
@@ -66,25 +66,25 @@ func ReadGolombUnsigned(buf []byte, pos *int) (uint32, error) {
 			break
 		}
 
-		leadingZeroBits++
-		if leadingZeroBits > 32 {
+		leadingZeros++
+		if leadingZeros > 32 {
 			return 0, fmt.Errorf("invalid value")
 		}
 	}
 
-	if (buflen*8 - *pos) < int(leadingZeroBits) {
+	if (buflen*8 - *pos) < int(leadingZeros) {
 		return 0, fmt.Errorf("not enough bits")
 	}
 
 	codeNum := uint32(0)
 
-	for n := leadingZeroBits; n > 0; n-- {
+	for n := leadingZeros; n > 0; n-- {
 		b := (buf[*pos>>0x03] >> (7 - (*pos & 0x07))) & 0x01
 		*pos++
 		codeNum |= uint32(b) << (n - 1)
 	}
 
-	codeNum = (1 << leadingZeroBits) - 1 + codeNum
+	codeNum = (1 << leadingZeros) - 1 + codeNum
 
 	return codeNum, nil
 }

diff --git a/pkg/codecs/av1/sequence_header.go b/pkg/codecs/av1/sequence_header.go
@@ -184,13 +184,43 @@ func (c *SequenceHeader_ColorConfig) unmarshal(seqProfile uint8, buf []byte, pos
 	return nil
 }
 
+// SequenceHeader_TimingInfo is the timing_info() struct in the AV1 specification.
+type SequenceHeader_TimingInfo struct { //nolint:revive
+	NumUnitsInDisplayTick    uint32
+	TimeScale                uint32
+	EqualPictureInterval     bool
+	NumTicksPerPictureMinus1 uint32
+}
+
+func (t *SequenceHeader_TimingInfo) unmarshal(buf []byte, pos *int) error {
+	err := bits.HasSpace(buf, *pos, 65)
+	if err != nil {
+		return err
+	}
+
+	t.NumUnitsInDisplayTick = uint32(bits.ReadBitsUnsafe(buf, pos, 32))
+	t.TimeScale = uint32(bits.ReadBitsUnsafe(buf, pos, 32))
+	t.EqualPictureInterval = bits.ReadFlagUnsafe(buf, pos)
+
+	if t.EqualPictureInterval {
+		t.NumTicksPerPictureMinus1, err = bits.ReadGolombUnsigned(buf, pos)
+		if err != nil {
+			return err
+		}
+	} else {
+		t.NumTicksPerPictureMinus1 = 0
+	}
+
+	return nil
+}
+
 // SequenceHeader is a AV1 Sequence header OBU.
 // Specification: https://aomediacodec.github.io/av1-spec/#sequence-header-obu-syntax
 type SequenceHeader struct {
 	SeqProfile                     uint8
 	StillPicture                   bool
 	ReducedStillPictureHeader      bool
-	TimingInfoPresentFlag          bool
+	TimingInfo                     *SequenceHeader_TimingInfo
 	DecoderModelInfoPresentFlag    bool
 	InitialDisplayDelayPresentFlag bool
 	OperatingPointsCntMinus1       uint8
@@ -203,6 +233,8 @@ type SequenceHeader struct {
 	MaxFrameWidthMinus1            uint32
 	MaxFrameHeightMinus1           uint32
 	FrameIDNumbersPresentFlag      bool
+	DeltaFrameIDLengthMinus2       uint8
+	AdditionalFrameIDLengthMinus1  uint8
 	Use128x128Superblock           bool
 	EnableFilterIntra              bool
 	EnableIntraEdgeFilter          bool
@@ -259,7 +291,7 @@ func (h *SequenceHeader) Unmarshal(buf []byte) error {
 	h.ReducedStillPictureHeader = bits.ReadFlagUnsafe(buf, &pos)
 
 	if h.ReducedStillPictureHeader {
-		h.TimingInfoPresentFlag = false
+		h.TimingInfo = nil
 		h.DecoderModelInfoPresentFlag = false
 		h.InitialDisplayDelayPresentFlag = false
 		h.OperatingPointsCntMinus1 = 0
@@ -275,15 +307,30 @@ func (h *SequenceHeader) Unmarshal(buf []byte) error {
 		h.DecoderModelPresentForThisOp = []bool{false}
 		h.InitialDisplayPresentForThisOp = []bool{false}
 	} else {
-		h.TimingInfoPresentFlag, err = bits.ReadFlag(buf, &pos)
+		var timingInfoPresentFlag bool
+		timingInfoPresentFlag, err = bits.ReadFlag(buf, &pos)
 		if err != nil {
 			return err
 		}
 
-		if h.TimingInfoPresentFlag {
-			return fmt.Errorf("timing_info_present_flag is not supported yet")
+		if timingInfoPresentFlag {
+			h.TimingInfo = &SequenceHeader_TimingInfo{}
+			err = h.TimingInfo.unmarshal(buf, &pos)
+			if err != nil {
+				return err
+			}
+
+			h.DecoderModelInfoPresentFlag, err = bits.ReadFlag(buf, &pos)
+			if err != nil {
+				return err
+			}
+			if h.DecoderModelInfoPresentFlag {
+				return fmt.Errorf("decoder_model_info_present_flag is not supported yet")
+			}
+		} else {
+			h.TimingInfo = nil
+			h.DecoderModelInfoPresentFlag = false
 		}
-		h.DecoderModelInfoPresentFlag = false
 
 		err = bits.HasSpace(buf, pos, 6)
 		if err != nil {
@@ -370,7 +417,13 @@ func (h *SequenceHeader) Unmarshal(buf []byte) error {
 		}
 
 		if h.FrameIDNumbersPresentFlag {
-			return fmt.Errorf("frame_id_numbers_present_flag is not supported yet")
+			err = bits.HasSpace(buf, pos, 7)
+			if err != nil {
+				return err
+			}
+
+			h.DeltaFrameIDLengthMinus2 = uint8(bits.ReadBitsUnsafe(buf, &pos, 4))
+			h.AdditionalFrameIDLengthMinus1 = uint8(bits.ReadBitsUnsafe(buf, &pos, 3))
 		}
 	}
 

diff --git a/pkg/codecs/av1/sequence_header_test.go b/pkg/codecs/av1/sequence_header_test.go
@@ -123,6 +123,50 @@ var casesSequenceHeader = []struct {
 		1920,
 		1080,
 	},
+	{
+		"amd hardware av1",
+		[]byte{
+			0x08, 0x04, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00,
+			0x00, 0xf3, 0x00, 0x00, 0x0e, 0x55, 0x77, 0xf8,
+			0x73, 0xd0, 0x02, 0x7d, 0x10, 0x10, 0x10, 0x10,
+			0x40,
+		},
+		SequenceHeader{
+			TimingInfo: &SequenceHeader_TimingInfo{
+				NumUnitsInDisplayTick: 1,
+				TimeScale:             60,
+				EqualPictureInterval:  true,
+			},
+			OperatingPointIdc:              []uint16{0},
+			SeqLevelIdx:                    []uint8{0x0e},
+			SeqTier:                        []bool{false},
+			DecoderModelPresentForThisOp:   []bool{false},
+			InitialDisplayPresentForThisOp: []bool{false},
+			InitialDisplayDelayMinus1:      []uint8{0},
+			MaxFrameWidthMinus1:            1919,
+			MaxFrameHeightMinus1:           1081,
+			FrameIDNumbersPresentFlag:      true,
+			DeltaFrameIDLengthMinus2:       13,
+			EnableOrderHint:                true,
+			SeqChooseScreenContentTools:    true,
+			SeqForceScreenContentTools:     SequenceHeader_SeqForceScreenContentTools(2),
+			SeqChooseIntegerMv:             true,
+			SeqForceIntegerMv:              2,
+			OrderHintBitsMinus1:            7,
+			EnableCdef:                     true,
+			ColorConfig: SequenceHeader_ColorConfig{
+				BitDepth:                    8,
+				ColorDescriptionPresentFlag: true,
+				ColorPrimaries:              1,
+				TransferCharacteristics:     SequenceHeader_TransferCharacteristics(1),
+				MatrixCoefficients:          SequenceHeader_MatrixCoefficients(1),
+				SubsamplingX:                true,
+				SubsamplingY:                true,
+			},
+		},
+		1920,
+		1082,
+	},
 }
 
 func TestSequenceHeaderUnmarshal(t *testing.T) {

diff --git a/pkg/codecs/av1/testdata/fuzz/FuzzSequenceHeaderUnmarshal/1d73d8e41097581e b/pkg/codecs/av1/testdata/fuzz/FuzzSequenceHeaderUnmarshal/1d73d8e41097581e
@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0700000007")
diff --git a/pkg/codecs/av1/testdata/fuzz/FuzzSequenceHeaderUnmarshal/8727b16d337d7b81 b/pkg/codecs/av1/testdata/fuzz/FuzzSequenceHeaderUnmarshal/8727b16d337d7b81
@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("00")
diff --git a/pkg/codecs/av1/testdata/fuzz/FuzzSequenceHeaderUnmarshal/95df00aea0ad3e95 b/pkg/codecs/av1/testdata/fuzz/FuzzSequenceHeaderUnmarshal/95df00aea0ad3e95
@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0700000000")
diff --git a/pkg/codecs/av1/testdata/fuzz/FuzzSequenceHeaderUnmarshal/a0866578e9bf0745 b/pkg/codecs/av1/testdata/fuzz/FuzzSequenceHeaderUnmarshal/a0866578e9bf0745
@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0700000002")
diff --git a/pkg/codecs/av1/testdata/fuzz/FuzzSequenceHeaderUnmarshal/fe4ea80246a7a881 b/pkg/codecs/av1/testdata/fuzz/FuzzSequenceHeaderUnmarshal/fe4ea80246a7a881
@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0700000001")