diff --git a/pkg/codecs/av1/bitstream.go b/pkg/codecs/av1/bitstream.go index 81a83a1..5bb2916 100644 --- a/pkg/codecs/av1/bitstream.go +++ b/pkg/codecs/av1/bitstream.go @@ -4,64 +4,69 @@ import ( "fmt" ) -func obuRemoveSize(h *OBUHeader, sizeN int, ob []byte) []byte { - newOBU := make([]byte, len(ob)-sizeN) - newOBU[0] = (byte(h.Type) << 3) - copy(newOBU[1:], ob[1+sizeN:]) - return newOBU -} - // BitstreamUnmarshal extracts a temporal unit from a bitstream. // Optionally, it also removes the size field from OBUs. +// +// Deprecated: replacted by Bitstream.Unmarshal. The removeSizeField has no effect anymore. +func BitstreamUnmarshal(buf []byte, _ bool) ([][]byte, error) { + var b Bitstream + err := b.Unmarshal(buf) + return b, err +} + +// BitstreamMarshal encodes a temporal unit into a bitstream. +// +// Deprecated: replacted by Bitstream.Marshal +func BitstreamMarshal(tu [][]byte) ([]byte, error) { + return Bitstream(tu).Marshal() +} + +// Bitstream is an AV1 bitstream. // Specification: https://aomediacodec.github.io/av1-spec/#low-overhead-bitstream-format -func BitstreamUnmarshal(bs []byte, removeSizeField bool) ([][]byte, error) { - var ret [][]byte +type Bitstream [][]byte +// Unmarshal decodes a Bitstream. +func (bs *Bitstream) Unmarshal(buf []byte) error { for { var h OBUHeader - err := h.Unmarshal(bs) + err := h.Unmarshal(buf) if err != nil { - return nil, err + return err } if !h.HasSize { - return nil, fmt.Errorf("OBU size not present") + return fmt.Errorf("OBU size not present") } var size LEB128 - n, err := size.Unmarshal(bs[1:]) + n, err := size.Unmarshal(buf[1:]) if err != nil { - return nil, err + return err } obuLen := 1 + n + int(size) - if len(bs) < obuLen { - return nil, fmt.Errorf("not enough bytes") + if len(buf) < obuLen { + return fmt.Errorf("not enough bytes") } - obu := bs[:obuLen] - - if removeSizeField { - obu = obuRemoveSize(&h, n, obu) - } + var obu []byte + obu, buf = buf[:obuLen], buf[obuLen:] - ret = append(ret, obu) - bs = bs[obuLen:] + *bs = append(*bs, obu) - if len(bs) == 0 { + if len(buf) == 0 { break } } - return ret, nil + return nil } -// BitstreamMarshal encodes a temporal unit into a bitstream. -// Specification: https://aomediacodec.github.io/av1-spec/#low-overhead-bitstream-format -func BitstreamMarshal(tu [][]byte) ([]byte, error) { +// Marshal encodes a Bitstream. +func (bs Bitstream) Marshal() ([]byte, error) { n := 0 - for _, obu := range tu { + for _, obu := range bs { n += len(obu) var h OBUHeader @@ -79,7 +84,7 @@ func BitstreamMarshal(tu [][]byte) ([]byte, error) { buf := make([]byte, n) n = 0 - for _, obu := range tu { + for _, obu := range bs { var h OBUHeader h.Unmarshal(obu) //nolint:errcheck @@ -89,6 +94,8 @@ func BitstreamMarshal(tu [][]byte) ([]byte, error) { size := len(obu) - 1 n += LEB128(uint32(size)).MarshalTo(buf[n:]) n += copy(buf[n:], obu[1:]) + } else { + n += copy(buf[n:], obu) } } diff --git a/pkg/codecs/av1/bitstream_test.go b/pkg/codecs/av1/bitstream_test.go index 4013863..972caf8 100644 --- a/pkg/codecs/av1/bitstream_test.go +++ b/pkg/codecs/av1/bitstream_test.go @@ -9,7 +9,7 @@ import ( var casesBitstream = []struct { name string enc []byte - dec [][]byte + dec Bitstream }{ { "standard", @@ -21,12 +21,12 @@ var casesBitstream = []struct { }, [][]byte{ { - 0x08, 0x00, 0x00, 0x00, 0x4a, 0xab, 0xbf, 0xc3, - 0x77, 0x6b, 0xe4, 0x40, 0x40, 0x40, 0x41, + 0x0a, 0x0e, 0x00, 0x00, 0x00, 0x4a, 0xab, 0xbf, + 0xc3, 0x77, 0x6b, 0xe4, 0x40, 0x40, 0x40, 0x41, }, { - 0x08, 0x00, 0x00, 0x00, 0x4a, 0xab, 0xbf, 0xc3, - 0x77, 0x6b, 0xe4, 0x40, 0x40, 0x40, 0x41, + 0x0a, 0x0e, 0x00, 0x00, 0x00, 0x4a, 0xab, 0xbf, + 0xc3, 0x77, 0x6b, 0xe4, 0x40, 0x40, 0x40, 0x41, }, }, }, @@ -35,7 +35,8 @@ var casesBitstream = []struct { func TestBitstreamUnmarshal(t *testing.T) { for _, ca := range casesBitstream { t.Run(ca.name, func(t *testing.T) { - dec, err := BitstreamUnmarshal(ca.enc, true) + var dec Bitstream + err := dec.Unmarshal(ca.enc) require.NoError(t, err) require.Equal(t, ca.dec, dec) }) @@ -45,7 +46,29 @@ func TestBitstreamUnmarshal(t *testing.T) { func TestBitstreamMarshal(t *testing.T) { for _, ca := range casesBitstream { t.Run(ca.name, func(t *testing.T) { - enc, err := BitstreamMarshal(ca.dec) + enc, err := ca.dec.Marshal() + require.NoError(t, err) + require.Equal(t, ca.enc, enc) + }) + } +} + +func TestBitstreamMarshalWithoutSize(t *testing.T) { + for _, ca := range casesBitstream { + t.Run(ca.name, func(t *testing.T) { + var tu Bitstream + for _, obu := range ca.dec { + var size LEB128 + n, err := size.Unmarshal(obu[1:]) + require.NoError(t, err) + + newObu := make([]byte, len(obu)-n) + newObu[0] = obu[0] & 0b01111000 + copy(newObu[1:], obu[1+n:]) + tu = append(tu, newObu) + } + + enc, err := tu.Marshal() require.NoError(t, err) require.Equal(t, ca.enc, enc) }) @@ -58,9 +81,10 @@ func FuzzBitstreamUnmarshal(f *testing.F) { } f.Fuzz(func(_ *testing.T, b []byte) { - tu, err := BitstreamUnmarshal(b, true) + var tu Bitstream + err := tu.Unmarshal(b) if err == nil { - BitstreamMarshal(tu) //nolint:errcheck + tu.Marshal() //nolint:errcheck } }) } diff --git a/pkg/formats/fmp4/init.go b/pkg/formats/fmp4/init.go index 93711bc..27a5103 100644 --- a/pkg/formats/fmp4/init.go +++ b/pkg/formats/fmp4/init.go @@ -26,8 +26,9 @@ const ( streamTypeAudioStream = 0x05 ) -func av1FindSequenceHeader(bs []byte) ([]byte, error) { - tu, err := av1.BitstreamUnmarshal(bs, true) +func av1FindSequenceHeader(buf []byte) ([]byte, error) { + var tu av1.Bitstream + err := tu.Unmarshal(buf) if err != nil { return nil, err } diff --git a/pkg/formats/fmp4/init_track.go b/pkg/formats/fmp4/init_track.go index 657aefe..ffffd84 100644 --- a/pkg/formats/fmp4/init_track.go +++ b/pkg/formats/fmp4/init_track.go @@ -333,8 +333,8 @@ func (it *InitTrack) marshal(w *mp4Writer) error { return err } - var bs []byte - bs, err = av1.BitstreamMarshal([][]byte{codec.SequenceHeader}) + var enc []byte + enc, err = av1.Bitstream([][]byte{codec.SequenceHeader}).Marshal() if err != nil { return err } @@ -351,7 +351,7 @@ func (it *InitTrack) marshal(w *mp4Writer) error { ChromaSubsamplingX: boolToUint8(av1SequenceHeader.ColorConfig.SubsamplingX), ChromaSubsamplingY: boolToUint8(av1SequenceHeader.ColorConfig.SubsamplingY), ChromaSamplePosition: uint8(av1SequenceHeader.ColorConfig.ChromaSamplePosition), - ConfigOBUs: bs, + ConfigOBUs: enc, }) if err != nil { return err diff --git a/pkg/formats/fmp4/part_sample.go b/pkg/formats/fmp4/part_sample.go index 086e827..6819e61 100644 --- a/pkg/formats/fmp4/part_sample.go +++ b/pkg/formats/fmp4/part_sample.go @@ -15,7 +15,7 @@ type PartSample struct { // NewPartSampleAV1 creates a sample with AV1 data. func NewPartSampleAV1(sequenceHeaderPresent bool, tu [][]byte) (*PartSample, error) { - bs, err := av1.BitstreamMarshal(tu) + bs, err := av1.Bitstream(tu).Marshal() if err != nil { return nil, err } @@ -42,7 +42,8 @@ func NewPartSampleH26x(ptsOffset int32, randomAccessPresent bool, au [][]byte) ( // GetAV1 gets AV1 data from the sample. func (ps PartSample) GetAV1() ([][]byte, error) { - tu, err := av1.BitstreamUnmarshal(ps.Payload, true) + var tu av1.Bitstream + err := tu.Unmarshal(ps.Payload) if err != nil { return nil, err } diff --git a/pkg/formats/pmp4/track.go b/pkg/formats/pmp4/track.go index 21a99be..f7d4ab5 100644 --- a/pkg/formats/pmp4/track.go +++ b/pkg/formats/pmp4/track.go @@ -355,7 +355,7 @@ func (t *Track) marshal(w *mp4Writer) (*headerTrackMarshalResult, error) { } var bs []byte - bs, err = av1.BitstreamMarshal([][]byte{codec.SequenceHeader}) + bs, err = av1.Bitstream([][]byte{codec.SequenceHeader}).Marshal() if err != nil { return nil, err }