Add an option to include SEI NALs for H264/H265 (#3313)

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

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"),

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

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