diff --git a/pkg/media/h264reader/h264reader.go b/pkg/media/h264reader/h264reader.go index 1b09e117..0c95ed83 100644 --- a/pkg/media/h264reader/h264reader.go +++ b/pkg/media/h264reader/h264reader.go @@ -18,6 +18,7 @@ type H264Reader struct { nalPrefixParsed bool readBuffer []byte tmpReadBuf []byte + includeSEI bool } var ( @@ -37,11 +38,42 @@ func NewReader(in io.Reader) (*H264Reader, error) { nalPrefixParsed: false, readBuffer: make([]byte, 0), tmpReadBuf: make([]byte, 4096), + includeSEI: false, } return reader, nil } +// Option configures the behavior of H264Reader. +type Option func(*H264Reader) error + +// NewReaderWithOptions creates new H264Reader with options. +// The default behavior is to skip SEI NAL units. +func NewReaderWithOptions(in io.Reader, options ...Option) (*H264Reader, error) { + reader, err := NewReader(in) + if err != nil { + return nil, err + } + + for _, option := range options { + if err := option(reader); err != nil { + return nil, err + } + } + + return reader, nil +} + +// WithIncludeSEI controls whether SEI (Supplemental Enhancement Information) NAL units are returned. +// Default is false (SEI is skipped). +func WithIncludeSEI(include bool) Option { + return func(r *H264Reader) error { + r.includeSEI = include + + return nil + } +} + // NAL H.264 Network Abstraction Layer. type NAL struct { PictureOrderCount uint32 @@ -145,7 +177,7 @@ func (reader *H264Reader) NextNAL() (*NAL, error) { if nalFound { nal := newNal(reader.nalBuffer) nal.parseHeader() - if nal.UnitType == NalUnitTypeSEI { + if !reader.includeSEI && nal.UnitType == NalUnitTypeSEI { reader.nalBuffer = nil continue diff --git a/pkg/media/h264reader/h264reader_test.go b/pkg/media/h264reader/h264reader_test.go index 9f6689e2..8090b4a3 100644 --- a/pkg/media/h264reader/h264reader_test.go +++ b/pkg/media/h264reader/h264reader_test.go @@ -102,6 +102,32 @@ func TestSkipSEI(t *testing.T) { require.Equal(byte(0xAB), nal.Data[0]) } +func TestIncludeSEI(t *testing.T) { + require := require.New(t) + h264Bytes := []byte{ + 0x0, 0x0, 0x0, 0x1, 0xAA, + 0x0, 0x0, 0x0, 0x1, 0x6, // SEI + 0x0, 0x0, 0x0, 0x1, 0xAB, + } + + reader, err := NewReaderWithOptions(bytes.NewReader(h264Bytes), WithIncludeSEI(true)) + require.NoError(err) + require.NotNil(reader) + + nal, err := reader.NextNAL() + require.NoError(err) + require.Equal(byte(0xAA), nal.Data[0]) + + nal, err = reader.NextNAL() + require.NoError(err) + require.Equal(NalUnitTypeSEI, nal.UnitType) + require.Equal(byte(0x6), nal.Data[0]) + + nal, err = reader.NextNAL() + require.NoError(err) + require.Equal(byte(0xAB), nal.Data[0]) +} + func TestIssue1734_NextNal(t *testing.T) { tt := [...][]byte{ []byte("\x00\x00\x010\x00\x00\x01\x00\x00\x01"), diff --git a/pkg/media/h265reader/h265reader.go b/pkg/media/h265reader/h265reader.go index 70e4ba51..6015f4c5 100644 --- a/pkg/media/h265reader/h265reader.go +++ b/pkg/media/h265reader/h265reader.go @@ -18,6 +18,7 @@ type H265Reader struct { nalPrefixParsed bool readBuffer []byte tmpReadBuf []byte + includeSEI bool } var ( @@ -25,6 +26,10 @@ var ( errDataIsNotH265Stream = errors.New("data is not a H265/HEVC bitstream") ) +func (reader *H265Reader) shouldSkipNAL(naluType NalUnitType) bool { + return !reader.includeSEI && (naluType == NalUnitTypePrefixSei || naluType == NalUnitTypeSuffixSei) +} + // NewReader creates new H265Reader. func NewReader(in io.Reader) (*H265Reader, error) { if in == nil { @@ -37,11 +42,42 @@ func NewReader(in io.Reader) (*H265Reader, error) { nalPrefixParsed: false, readBuffer: make([]byte, 0), tmpReadBuf: make([]byte, 4096), + includeSEI: false, } return reader, nil } +// Option configures the behavior of H265Reader. +type Option func(*H265Reader) error + +// NewReaderWithOptions creates new H265Reader with options. +// The default behavior is to skip SEI NAL units. +func NewReaderWithOptions(in io.Reader, options ...Option) (*H265Reader, error) { + reader, err := NewReader(in) + if err != nil { + return nil, err + } + + for _, option := range options { + if err := option(reader); err != nil { + return nil, err + } + } + + return reader, nil +} + +// WithIncludeSEI controls whether SEI (Supplemental Enhancement Information) NAL units are returned. +// Default is false (SEI is skipped). +func WithIncludeSEI(include bool) Option { + return func(r *H265Reader) error { + r.includeSEI = include + + return nil + } +} + // NAL H.265/HEVC Network Abstraction Layer. type NAL struct { PictureOrderCount uint32 @@ -151,7 +187,7 @@ func (reader *H265Reader) NextNAL() (*NAL, error) { nalFound := reader.processByte(readByte) if nalFound { naluType := NalUnitType((reader.nalBuffer[0] & 0x7E) >> 1) - if naluType == NalUnitTypePrefixSei || naluType == NalUnitTypeSuffixSei { + if reader.shouldSkipNAL(naluType) { reader.nalBuffer = nil continue diff --git a/pkg/media/h265reader/h265reader_test.go b/pkg/media/h265reader/h265reader_test.go index 7258ece0..0f4f6f29 100644 --- a/pkg/media/h265reader/h265reader_test.go +++ b/pkg/media/h265reader/h265reader_test.go @@ -90,6 +90,74 @@ func TestH265Reader_processByte(t *testing.T) { assert.Equal(t, 0, reader.countOfConsecutiveZeroBytes) } +func TestH265Reader_SEIPrefixSuffixSkippedByDefault(t *testing.T) { + // Build a small Annex-B stream: VPS, PrefixSEI, SuffixSEI, SPS + // NAL header is 2 bytes. NAL unit type is encoded in bits 1..6 of the first byte. + stream := []byte{ + 0x0, 0x0, 0x0, 0x1, 0x40, 0x01, 0xFF, // VPS (type 32) + 0x0, 0x0, 0x0, 0x1, 0x4E, 0x01, 0xFF, // PrefixSEI (type 39) + 0x0, 0x0, 0x0, 0x1, 0x50, 0x01, 0xFF, // SuffixSEI (type 40) + 0x0, 0x0, 0x0, 0x1, 0x42, 0x01, 0xFF, // SPS (type 33) + } + + reader, err := NewReader(bytes.NewReader(stream)) + assert.NoError(t, err) + + nal1, err := reader.NextNAL() + assert.NoError(t, err) + assert.NotNil(t, nal1) + assert.Equal(t, NalUnitTypeVps, nal1.NalUnitType) + + // SEI should be skipped by default (both Prefix and Suffix) + nal2, err := reader.NextNAL() + assert.NoError(t, err) + assert.NotNil(t, nal2) + assert.Equal(t, NalUnitTypeSps, nal2.NalUnitType) + + _, err = reader.NextNAL() + assert.Equal(t, io.EOF, err) +} + +func TestH265Reader_IncludeSEI(t *testing.T) { + vps := []byte{0x40, 0x01} // NalUnitTypeVps (32) + prefixSEI := []byte{0x4E, 0x01} // NalUnitTypePrefixSei (39) + suffixSEI := []byte{0x50, 0x01} // NalUnitTypeSuffixSei (40) + sps := []byte{0x42, 0x01} // NalUnitTypeSps (33) + start := []byte{0x0, 0x0, 0x0, 0x1} + + stream := append([]byte{}, start...) + stream = append(stream, vps...) + stream = append(stream, start...) + stream = append(stream, prefixSEI...) + stream = append(stream, start...) + stream = append(stream, suffixSEI...) + stream = append(stream, start...) + stream = append(stream, sps...) + + reader, err := NewReaderWithOptions(bytes.NewReader(stream), WithIncludeSEI(true)) + assert.NoError(t, err) + + nal1, err := reader.NextNAL() + assert.NoError(t, err) + assert.NotNil(t, nal1) + assert.Equal(t, NalUnitTypeVps, nal1.NalUnitType) + + nal2, err := reader.NextNAL() + assert.NoError(t, err) + assert.NotNil(t, nal2) + assert.Equal(t, NalUnitTypePrefixSei, nal2.NalUnitType) + + nal3, err := reader.NextNAL() + assert.NoError(t, err) + assert.NotNil(t, nal3) + assert.Equal(t, NalUnitTypeSuffixSei, nal3.NalUnitType) + + nal4, err := reader.NextNAL() + assert.NoError(t, err) + assert.NotNil(t, nal4) + assert.Equal(t, NalUnitTypeSps, nal4.NalUnitType) +} + func TestNAL_parseHeader(t *testing.T) { // Test VPS NAL header parsing data := []byte{0x40, 0x01, 0x0C, 0x01} // VPS NAL unit